feat(image): return error early if total size of layers exceeds limit (#8294)

Signed-off-by: nikpivkin <[email protected]>
Signed-off-by: knqyf263 <[email protected]>
Co-authored-by: knqyf263 <[email protected]>

Author: nikpivkin

docs/docs/target/container_image.md

@@ -526,7 +526,14 @@ $ trivy image --podman-host /run/user/1000/podman/podman.sock YOUR_IMAGE
 ```
 
 ### Prevent scanning oversized container images
-Use the `--max-image-size` flag to avoid scanning images that exceed a specified size. The size is specified in a human-readable format (e.g., `100MB`, `10GB`). If the compressed image size exceeds the specified threshold, an error is returned immediately. Otherwise, all layers are pulled, stored in a temporary folder, and their uncompressed size is verified before scanning. Temporary layers are always cleaned up, even after a successful scan.
+Use the `--max-image-size` flag to avoid scanning images that exceed a specified size. The size is specified in a human-readable format[^1] (e.g., `100MB`, `10GB`).
+
+An error is returned in the following cases:
+
+- if the compressed image size exceeds the limit,
+- if the accumulated size of the uncompressed layers exceeds the limit during their pulling.
+
+The layers are pulled into a temporary folder during their pulling and are always cleaned up, even after a successful scan.
 
 !!! warning "EXPERIMENTAL"
     This feature might change without preserving backwards compatibility.
@@ -542,3 +549,5 @@ Error Output:
 ```bash
 Error: uncompressed image size (15GB) exceeds maximum allowed size (10GB)
 ```
+
+[^1]: Trivy uses decimal (SI) prefixes (based on 1000) for size.

integration/docker_engine_test.go

@@ -206,7 +206,7 @@ func TestDockerEngine(t *testing.T) {
 			name:         "sad path, image size is larger than the maximum",
 			input:        "testdata/fixtures/images/alpine-39.tar.gz",
 			maxImageSize: "3mb",
-			wantErr:      "uncompressed image size 5.8MB exceeds maximum allowed size 3MB",
+			wantErr:      "uncompressed layers size 5.8MB exceeds maximum allowed size 3MB",
 		},
 	}
 

pkg/fanal/artifact/image/image.go

@@ -219,68 +219,60 @@ func (a Artifact) consolidateCreatedBy(diffIDs, layerKeys []string, configFile *
 	return layerKeyMap
 }
 
-func limitErrorMessage(typ string, maxSize, imageSize int64) string {
-	return fmt.Sprintf(
-		"%s image size %s exceeds maximum allowed size %s", typ,
-		units.HumanSizeWithPrecision(float64(imageSize), 3),
-		units.HumanSize(float64(maxSize)),
-	)
+func (a Artifact) imageSizeError(typ string, size int64) error {
+	return &trivyTypes.UserError{
+		Message: fmt.Sprintf(
+			"%s size %s exceeds maximum allowed size %s", typ,
+			units.HumanSizeWithPrecision(float64(size), 3),
+			units.HumanSize(float64(a.artifactOption.ImageOption.MaxImageSize)),
+		),
+	}
 }
 
 func (a Artifact) checkImageSize(ctx context.Context, diffIDs []string) error {
-	maxSize := a.artifactOption.ImageOption.MaxImageSize
-	if maxSize == 0 {
+	if a.artifactOption.ImageOption.MaxImageSize == 0 {
 		return nil
 	}
 
-	compressedSize, err := a.compressedImageSize(diffIDs)
-	if err != nil {
+	if err := a.checkCompressedImageSize(diffIDs); err != nil {
 		return xerrors.Errorf("failed to get compressed image size: %w", err)
 	}
 
-	if compressedSize > maxSize {
-		return &trivyTypes.UserError{
-			Message: limitErrorMessage("compressed", maxSize, compressedSize),
-		}
-	}
-
-	imageSize, err := a.imageSize(ctx, diffIDs)
-	if err != nil {
+	if err := a.checkUncompressedImageSize(ctx, diffIDs); err != nil {
 		return xerrors.Errorf("failed to calculate image size: %w", err)
 	}
-
-	if imageSize > maxSize {
-		return &trivyTypes.UserError{
-			Message: limitErrorMessage("uncompressed", maxSize, imageSize),
-		}
-	}
 	return nil
 }
 
-func (a Artifact) compressedImageSize(diffIDs []string) (int64, error) {
+func (a Artifact) checkCompressedImageSize(diffIDs []string) error {
 	var totalSize int64
+
 	for _, diffID := range diffIDs {
 		h, err := v1.NewHash(diffID)
 		if err != nil {
-			return -1, xerrors.Errorf("invalid layer ID (%s): %w", diffID, err)
+			return xerrors.Errorf("invalid layer ID (%s): %w", diffID, err)
 		}
 
 		layer, err := a.image.LayerByDiffID(h)
 		if err != nil {
-			return -1, xerrors.Errorf("failed to get the layer (%s): %w", diffID, err)
+			return xerrors.Errorf("failed to get the layer (%s): %w", diffID, err)
 		}
 		layerSize, err := layer.Size()
 		if err != nil {
-			return -1, xerrors.Errorf("failed to get layer size: %w", err)
+			return xerrors.Errorf("failed to get layer size: %w", err)
 		}
 		totalSize += layerSize
 	}
 
-	return totalSize, nil
+	if totalSize > a.artifactOption.ImageOption.MaxImageSize {
+		return a.imageSizeError("compressed image", totalSize)
+	}
+
+	return nil
 }
 
-func (a Artifact) imageSize(ctx context.Context, diffIDs []string) (int64, error) {
-	var imageSize int64
+func (a Artifact) checkUncompressedImageSize(ctx context.Context, diffIDs []string) error {
+	var totalSize int64
 
 	p := parallel.NewPipeline(a.artifactOption.Parallel, false, diffIDs,
 		func(_ context.Context, diffID string) (int64, error) {
@@ -291,16 +283,19 @@ func (a Artifact) imageSize(ctx context.Context, diffIDs []string) (int64, error
 			return layerSize, nil
 		},
 		func(layerSize int64) error {
-			imageSize += layerSize
+			totalSize += layerSize
+			if totalSize > a.artifactOption.ImageOption.MaxImageSize {
+				return a.imageSizeError("uncompressed layers", totalSize)
+			}
 			return nil
 		},
 	)
 
 	if err := p.Do(ctx); err != nil {
-		return -1, xerrors.Errorf("pipeline error: %w", err)
+		return xerrors.Errorf("pipeline error: %w", err)
 	}
 
-	return imageSize, nil
+	return nil
 }
 
 func (a Artifact) saveLayer(diffID string) (int64, error) {

pkg/fanal/artifact/image/image_test.go

@@ -3,11 +3,13 @@ package image_test
 import (
 	"context"
 	"errors"
+	"math/rand"
 	"testing"
 	"time"
 
 	"github.com/docker/go-units"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/random"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/xerrors"
@@ -2245,22 +2247,6 @@ func TestArtifact_Inspect(t *testing.T) {
 			},
 			wantErr: "put artifact failed",
 		},
-		{
-			name:      "sad path, compressed image size is larger than the maximum",
-			imagePath: "../../test/testdata/alpine-311.tar.gz",
-			artifactOpt: artifact.Option{
-				ImageOption: types.ImageOptions{MaxImageSize: units.MB * 1},
-			},
-			wantErr: "compressed image size 3.03MB exceeds maximum allowed size 1MB",
-		},
-		{
-			name:      "sad path, image size is larger than the maximum",
-			imagePath: "../../test/testdata/alpine-311.tar.gz",
-			artifactOpt: artifact.Option{
-				ImageOption: types.ImageOptions{MaxImageSize: units.MB * 4},
-			},
-			wantErr: "uncompressed image size 5.86MB exceeds maximum allowed size 4MB",
-		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -2287,3 +2273,42 @@ func TestArtifact_Inspect(t *testing.T) {
 		})
 	}
 }
+
+func TestArtifact_InspectWithMaxImageSize(t *testing.T) {
+	randomImage, err := random.Image(1000, 2, random.WithSource(rand.NewSource(0)))
+	require.NoError(t, err)
+
+	img := &fakeImage{Image: randomImage}
+	mockCache := new(cache.MockArtifactCache)
+
+	tests := []struct {
+		name        string
+		artifactOpt artifact.Option
+		wantErr     string
+	}{
+		{
+			name: "compressed image size is larger than the maximum",
+			artifactOpt: artifact.Option{
+				ImageOption: types.ImageOptions{MaxImageSize: units.KB * 1},
+			},
+			wantErr: "compressed image size 2.44kB exceeds maximum allowed size 1kB",
+		},
+		{
+			name: "uncompressed layers size is larger than the maximum",
+			artifactOpt: artifact.Option{
+				ImageOption: types.ImageOptions{MaxImageSize: units.KB * 3},
+			},
+			wantErr: "uncompressed layers size 5.12kB exceeds maximum allowed size 3kB",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			artifact, err := image2.NewArtifact(img, mockCache, tt.artifactOpt)
+			require.NoError(t, err)
+
+			_, err = artifact.Inspect(context.Background())
+			require.ErrorContains(t, err, tt.wantErr)
+		})
+	}
+}

pkg/fanal/artifact/image/remote_sbom_test.go

@@ -31,7 +31,7 @@ func TestMain(m *testing.M) {
 type fakeImage struct {
 	name        string
 	repoDigests []string
-	*fakei.FakeImage
+	v1.Image
 	types.ImageExtension
 }
 
@@ -160,7 +160,7 @@ func TestArtifact_InspectRekorAttestation(t *testing.T) {
 			img := &fakeImage{
 				name:        tt.fields.imageName,
 				repoDigests: tt.fields.repoDigests,
-				FakeImage:   fi,
+				Image:       fi,
 			}
 			a, err := image2.NewArtifact(img, mockCache, tt.artifactOpt)
 			require.NoError(t, err)
@@ -304,7 +304,7 @@ func TestArtifact_inspectOCIReferrerSBOM(t *testing.T) {
 			img := &fakeImage{
 				name:        tt.fields.imageName,
 				repoDigests: tt.fields.repoDigests,
-				FakeImage:   fi,
+				Image:       fi,
 			}
 			a, err := image2.NewArtifact(img, mockCache, tt.artifactOpt)
 			require.NoError(t, err)