Fix excessive memory and disk usage when sharing layers

storage/storage_dest.go

@@ -21,7 +21,6 @@ import (
 	"github.com/containers/image/v5/internal/imagedestination/stubs"
 	"github.com/containers/image/v5/internal/private"
 	"github.com/containers/image/v5/internal/putblobdigest"
-	"github.com/containers/image/v5/internal/set"
 	"github.com/containers/image/v5/internal/signature"
 	"github.com/containers/image/v5/internal/tmpdir"
 	"github.com/containers/image/v5/manifest"
@@ -121,6 +120,9 @@ type storageImageDestinationLockProtected struct {
 	filenames map[digest.Digest]string
 	// Mapping from layer blobsums to their sizes. If set, filenames and blobDiffIDs must also be set.
 	fileSizes map[digest.Digest]int64
+
+	// Config
+	configDigest digest.Digest // "" if N/A or not known yet.
 }
 
 // addedLayerInfo records data about a layer to use in this image.
@@ -214,7 +216,17 @@ func (s *storageImageDestination) PutBlobWithOptions(ctx context.Context, stream
 		return info, err
 	}
 
-	if options.IsConfig || options.LayerIndex == nil {
+	if options.IsConfig {
+		s.lock.Lock()
+		defer s.lock.Unlock()
+		if s.lockProtected.configDigest != "" {
+			return private.UploadedBlob{}, fmt.Errorf("after config %q, refusing to record another config %q",
+				s.lockProtected.configDigest.String(), info.Digest.String())
+		}
+		s.lockProtected.configDigest = info.Digest
+		return info, nil
+	}
+	if options.LayerIndex == nil {
 		return info, nil
 	}
 
@@ -1210,22 +1222,14 @@ func (s *storageImageDestination) CommitWithOptions(ctx context.Context, options
 		imgOptions.CreationDate = *inspect.Created
 	}
 
-	// Set up to save the non-layer blobs as data items.  Since we only share layers, they should all be in files, so
-	// we just need to screen out the ones that are actually layers to get the list of non-layers.
-	dataBlobs := set.New[digest.Digest]()
-	for blob := range s.lockProtected.filenames {
-		dataBlobs.Add(blob)
-	}
-	for _, layerBlob := range layerBlobs {
-		dataBlobs.Delete(layerBlob.Digest)
-	}
-	for _, blob := range dataBlobs.Values() {
-		v, err := os.ReadFile(s.lockProtected.filenames[blob])
+	// Set up to save the config as a data item.  Since we only share layers, the config should be in a file.
+	if s.lockProtected.configDigest != "" {
+		v, err := os.ReadFile(s.lockProtected.filenames[s.lockProtected.configDigest])
 		if err != nil {
-			return fmt.Errorf("copying non-layer blob %q to image: %w", blob, err)
+			return fmt.Errorf("copying config blob %q to image: %w", s.lockProtected.configDigest, err)
 		}
 		imgOptions.BigData = append(imgOptions.BigData, storage.ImageBigDataOption{
-			Key:    blob.String(),
+			Key:    s.lockProtected.configDigest.String(),
 			Data:   v,
 			Digest: digest.Canonical.FromBytes(v),
 		})

storage/storage_test.go

@@ -277,11 +277,11 @@ func makeLayer(t *testing.T, compression archive.Compression) testBlob {
 	}
 }
 
-func (l testBlob) storeBlob(t *testing.T, dest types.ImageDestination, cache types.BlobInfoCache, mimeType string) manifest.Schema2Descriptor {
+func (l testBlob) storeBlob(t *testing.T, dest types.ImageDestination, cache types.BlobInfoCache, mimeType string, isConfig bool) manifest.Schema2Descriptor {
 	_, err := dest.PutBlob(context.Background(), bytes.NewReader(l.data), types.BlobInfo{
 		Size:   l.compressedSize,
 		Digest: l.compressedDigest,
-	}, cache, false)
+	}, cache, isConfig)
 	require.NoError(t, err)
 	return manifest.Schema2Descriptor{
 		MediaType: mimeType,
@@ -312,13 +312,13 @@ func createUncommittedImageDest(t *testing.T, ref types.ImageReference, cache ty
 
 	layerDescriptors := []manifest.Schema2Descriptor{}
 	for _, layer := range layers {
-		desc := layer.storeBlob(t, dest, cache, manifest.DockerV2Schema2LayerMediaType)
+		desc := layer.storeBlob(t, dest, cache, manifest.DockerV2Schema2LayerMediaType, false)
 		layerDescriptors = append(layerDescriptors, desc)
 	}
 
 	var configDescriptor manifest.Schema2Descriptor
 	if config != nil {
-		configDescriptor = config.storeBlob(t, dest, cache, manifest.DockerV2Schema2ConfigMediaType)
+		configDescriptor = config.storeBlob(t, dest, cache, manifest.DockerV2Schema2ConfigMediaType, true)
 	} else {
 		// Use a random digest so that different calls to createUncommittedImageDest with config == nil don’t try to
 		// use the same image ID.
@@ -441,9 +441,9 @@ func TestWriteRead(t *testing.T) {
 			compression = archive.Gzip
 		}
 		layer := makeLayer(t, compression)
-		_ = layer.storeBlob(t, dest, cache, manifest.DockerV2Schema2LayerMediaType)
+		_ = layer.storeBlob(t, dest, cache, manifest.DockerV2Schema2LayerMediaType, false)
 		t.Logf("Wrote randomly-generated layer %q (%d/%d bytes) to destination", layer.compressedDigest, layer.compressedSize, layer.uncompressedSize)
-		_ = config.storeBlob(t, dest, cache, manifest.DockerV2Schema2ConfigMediaType)
+		_ = config.storeBlob(t, dest, cache, manifest.DockerV2Schema2ConfigMediaType, true)
 
 		manifest := strings.ReplaceAll(manifestFmt, "%lh", layer.compressedDigest.String())
 		manifest = strings.ReplaceAll(manifest, "%ch", config.compressedDigest.String())