feat(image): prevent scanning oversized container images (#8178)

Signed-off-by: nikpivkin <nikita.pivkin@smartforce.io>
Signed-off-by: knqyf263 <knqyf263@gmail.com>
Co-authored-by: knqyf263 <knqyf263@gmail.com>

cmd/trivy/main.go

@@ -21,6 +21,12 @@ func main() {
 		if errors.As(err, &exitError) {
 			os.Exit(exitError.Code)
 		}
+
+		var userErr *types.UserError
+		if errors.As(err, &userErr) {
+			log.Fatal("Error", log.Err(userErr))
+		}
+
 		log.Fatal("Fatal error", log.Err(err))
 	}
 }

docs/docs/references/configuration/cli/trivy_image.md

@@ -79,6 +79,7 @@ trivy image [flags] IMAGE_NAME
       --license-confidence-level float    specify license classifier's confidence level (default 0.9)
       --license-full                      eagerly look for licenses in source code headers and license files
       --list-all-pkgs                     output all packages in the JSON report regardless of vulnerability
+      --max-image-size string             [EXPERIMENTAL] maximum image size to process, specified in a human-readable format (e.g., '44kB', '17MB'); an error will be returned if the image exceeds this size
       --misconfig-scanners strings        comma-separated list of misconfig scanners to use for misconfiguration scanning (default [azure-arm,cloudformation,dockerfile,helm,kubernetes,terraform,terraformplan-json,terraformplan-snapshot])
       --module-dir string                 specify directory to the wasm modules that will be loaded (default "$HOME/.trivy/modules")
       --no-progress                       suppress progress bar

docs/docs/references/configuration/config-file.md

@@ -137,6 +137,9 @@ image:
   # Same as '--input'
   input: ""
 
+  # Same as '--max-image-size'
+  max-size: ""
+
   # Same as '--platform'
   platform: ""
 

docs/docs/target/container_image.md

@@ -518,3 +518,21 @@ You can configure Podman daemon socket with `--podman-host`.
 ```shell
 $ 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.
+
+!!! warning "EXPERIMENTAL"
+    This feature might change without preserving backwards compatibility.
+
+
+Example Usage:
+```bash
+# Limit uncompressed image size to 10GB
+$ trivy image --max-image-size=10GB myapp:latest
+```
+
+Error Output:
+```bash
+Error: uncompressed image size (15GB) exceeds maximum allowed size (10GB)
+```

go.mod

@@ -46,6 +46,7 @@ require (
 	github.com/docker/cli v27.5.0+incompatible
 	github.com/docker/docker v27.5.0+incompatible
 	github.com/docker/go-connections v0.5.0
+	github.com/docker/go-units v0.5.0
 	github.com/fatih/color v1.18.0
 	github.com/go-git/go-git/v5 v5.13.1
 	github.com/go-openapi/runtime v0.28.0 // indirect
@@ -222,7 +223,6 @@ require (
 	github.com/docker/distribution v2.8.3+incompatible // indirect
 	github.com/docker/docker-credential-helpers v0.8.2 // indirect
 	github.com/docker/go-metrics v0.0.1 // indirect
-	github.com/docker/go-units v0.5.0 // indirect
 	github.com/docker/libtrust v0.0.0-20160708172513-aabc10ec26b7 // indirect
 	github.com/dsnet/compress v0.0.1 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect

integration/docker_engine_test.go

@@ -25,6 +25,7 @@ func TestDockerEngine(t *testing.T) {
 		ignoreStatus  []string
 		severity      []string
 		ignoreIDs     []string
+		maxImageSize  string
 		input         string
 		golden        string
 		wantErr       string
@@ -34,6 +35,12 @@ func TestDockerEngine(t *testing.T) {
 			input:  "testdata/fixtures/images/alpine-39.tar.gz",
 			golden: "testdata/alpine-39.json.golden",
 		},
+		{
+			name:         "alpine:3.9, with max image size",
+			maxImageSize: "100mb",
+			input:        "testdata/fixtures/images/alpine-39.tar.gz",
+			golden:       "testdata/alpine-39.json.golden",
+		},
 		{
 			name: "alpine:3.9, with high and critical severity",
 			severity: []string{
@@ -195,6 +202,12 @@ func TestDockerEngine(t *testing.T) {
 			input:        "badimage:latest",
 			wantErr:      "unable to inspect the image (badimage:latest)",
 		},
+		{
+			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",
+		},
 	}
 
 	// Set up testing DB
@@ -263,6 +276,11 @@ func TestDockerEngine(t *testing.T) {
 				require.NoError(t, err, "failed to write .trivyignore")
 				defer os.Remove(trivyIgnore)
 			}
+
+			if tt.maxImageSize != "" {
+				osArgs = append(osArgs, []string{"--max-image-size", tt.maxImageSize}...)
+			}
+
 			osArgs = append(osArgs, tt.input)
 
 			// Run Trivy

pkg/commands/artifact/run.go

@@ -587,6 +587,7 @@ func (r *runner) initScannerConfig(ctx context.Context, opts flag.Options) (Scan
 					Host: opts.PodmanHost,
 				},
 				ImageSources: opts.ImageSources,
+				MaxImageSize: opts.MaxImageSize,
 			},
 
 			// For misconfiguration scanning

pkg/fanal/artifact/image/image.go

@@ -3,13 +3,16 @@ package image
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"os"
+	"path/filepath"
 	"reflect"
 	"slices"
 	"strings"
 	"sync"
 
+	"github.com/docker/go-units"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/samber/lo"
 	"golang.org/x/xerrors"
@@ -24,6 +27,7 @@ import (
 	"github.com/aquasecurity/trivy/pkg/log"
 	"github.com/aquasecurity/trivy/pkg/parallel"
 	"github.com/aquasecurity/trivy/pkg/semaphore"
+	trivyTypes "github.com/aquasecurity/trivy/pkg/types"
 )
 
 type Artifact struct {
@@ -36,6 +40,8 @@ type Artifact struct {
 	handlerManager handler.Manager
 
 	artifactOption artifact.Option
+
+	layerCacheDir string
 }
 
 type LayerInfo struct {
@@ -60,6 +66,11 @@ func NewArtifact(img types.Image, c cache.ArtifactCache, opt artifact.Option) (a
 		return nil, xerrors.Errorf("config analyzer group error: %w", err)
 	}
 
+	cacheDir, err := os.MkdirTemp("", "layers")
+	if err != nil {
+		return nil, xerrors.Errorf("failed to create a cache layers temp dir: %w", err)
+	}
+
 	return Artifact{
 		logger:         log.WithPrefix("image"),
 		image:          img,
@@ -70,10 +81,11 @@ func NewArtifact(img types.Image, c cache.ArtifactCache, opt artifact.Option) (a
 		handlerManager: handlerManager,
 
 		artifactOption: opt,
+		layerCacheDir:  cacheDir,
 	}, nil
 }
 
-func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) {
+func (a Artifact) Inspect(ctx context.Context) (ref artifact.Reference, err error) {
 	imageID, err := a.image.ID()
 	if err != nil {
 		return artifact.Reference{}, xerrors.Errorf("unable to get the image ID: %w", err)
@@ -88,6 +100,15 @@ func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) {
 	diffIDs := a.diffIDs(configFile)
 	a.logger.Debug("Detected diff ID", log.Any("diff_ids", diffIDs))
 
+	defer func() {
+		if rerr := os.RemoveAll(a.layerCacheDir); rerr != nil {
+			log.Error("Failed to remove layer cache", log.Err(rerr))
+		}
+	}()
+	if err := a.checkImageSize(ctx, diffIDs); err != nil {
+		return artifact.Reference{}, err
+	}
+
 	// Try retrieving a remote SBOM document
 	if res, err := a.retrieveRemoteSBOM(ctx); err == nil {
 		// Found SBOM
@@ -141,7 +162,7 @@ func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) {
 	}, nil
 }
 
-func (Artifact) Clean(_ artifact.Reference) error {
+func (a Artifact) Clean(_ artifact.Reference) error {
 	return nil
 }
 
@@ -198,6 +219,107 @@ 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) checkImageSize(ctx context.Context, diffIDs []string) error {
+	maxSize := a.artifactOption.ImageOption.MaxImageSize
+	if maxSize == 0 {
+		return nil
+	}
+
+	compressedSize, err := a.compressedImageSize(diffIDs)
+	if 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 {
+		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) {
+	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)
+		}
+
+		layer, err := a.image.LayerByDiffID(h)
+		if err != nil {
+			return -1, 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)
+		}
+		totalSize += layerSize
+	}
+
+	return totalSize, nil
+}
+
+func (a Artifact) imageSize(ctx context.Context, diffIDs []string) (int64, error) {
+	var imageSize int64
+
+	p := parallel.NewPipeline(a.artifactOption.Parallel, false, diffIDs,
+		func(_ context.Context, diffID string) (int64, error) {
+			layerSize, err := a.saveLayer(diffID)
+			if err != nil {
+				return -1, xerrors.Errorf("failed to save layer: %w", err)
+			}
+			return layerSize, nil
+		},
+		func(layerSize int64) error {
+			imageSize += layerSize
+			return nil
+		},
+	)
+
+	if err := p.Do(ctx); err != nil {
+		return -1, xerrors.Errorf("pipeline error: %w", err)
+	}
+
+	return imageSize, nil
+}
+
+func (a Artifact) saveLayer(diffID string) (int64, error) {
+	a.logger.Debug("Pulling the layer to the local cache", log.String("diff_id", diffID))
+	_, rc, err := a.uncompressedLayer(diffID)
+	if err != nil {
+		return -1, xerrors.Errorf("unable to get uncompressed layer %s: %w", diffID, err)
+	}
+	defer rc.Close()
+
+	f, err := os.Create(filepath.Join(a.layerCacheDir, diffID))
+	if err != nil {
+		return -1, xerrors.Errorf("failed to create a file: %w", err)
+	}
+	defer f.Close()
+
+	return io.Copy(f, rc)
+}
+
 func (a Artifact) inspect(ctx context.Context, missingImage string, layerKeys, baseDiffIDs []string,
 	layerKeyMap map[string]LayerInfo, configFile *v1.ConfigFile) error {
 
@@ -361,6 +483,12 @@ func (a Artifact) uncompressedLayer(diffID string) (string, io.ReadCloser, error
 		digest = d.String()
 	}
 
+	f, err := os.Open(filepath.Join(a.layerCacheDir, diffID))
+	if err == nil {
+		a.logger.Debug("Loaded the layer from the local cache", log.String("diff_id", diffID))
+		return digest, f, nil
+	}
+
 	rc, err := layer.Uncompressed()
 	if err != nil {
 		return "", nil, xerrors.Errorf("failed to get the layer content (%s): %w", diffID, err)

pkg/fanal/artifact/image/image_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/docker/go-units"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -348,6 +349,7 @@ func TestArtifact_Inspect(t *testing.T) {
 			imagePath: "../../test/testdata/alpine-311.tar.gz",
 			artifactOpt: artifact.Option{
 				LicenseScannerOption: analyzer.LicenseScannerOption{Full: true},
+				ImageOption:          types.ImageOptions{MaxImageSize: units.GB},
 			},
 			missingBlobsExpectation: cache.ArtifactCacheMissingBlobsExpectation{
 				Args: cache.ArtifactCacheMissingBlobsArgs{
@@ -2243,6 +2245,22 @@ 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) {
@@ -2262,6 +2280,8 @@ func TestArtifact_Inspect(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr, tt.name)
 				return
 			}
+			defer a.Clean(got)
+
 			require.NoError(t, err, tt.name)
 			assert.Equal(t, tt.want, got)
 		})

pkg/fanal/artifact/image/remote_sbom_test.go

@@ -170,6 +170,8 @@ func TestArtifact_InspectRekorAttestation(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr)
 				return
 			}
+			defer a.Clean(got)
+
 			require.NoError(t, err, tt.name)
 			got.BOM = nil
 			assert.Equal(t, tt.want, got)
@@ -312,6 +314,7 @@ func TestArtifact_inspectOCIReferrerSBOM(t *testing.T) {
 				assert.ErrorContains(t, err, tt.wantErr)
 				return
 			}
+			defer a.Clean(got)
 
 			require.NoError(t, err, tt.name)
 			got.BOM = nil

pkg/fanal/test/integration/registry_test.go

@@ -256,6 +256,7 @@ func analyze(ctx context.Context, imageRef string, opt types.ImageOptions) (*typ
 	if err != nil {
 		return nil, err
 	}
+	defer ar.Clean(imageInfo)
 
 	imageDetail, err := ap.ApplyLayers(imageInfo.ID, imageInfo.BlobIDs)
 	if err != nil {

pkg/fanal/types/image.go

@@ -53,6 +53,7 @@ type ImageOptions struct {
 	PodmanOptions     PodmanOptions
 	ContainerdOptions ContainerdOptions
 	ImageSources      ImageSources
+	MaxImageSize      int64
 }
 
 type DockerOptions struct {

pkg/flag/image_flags.go

@@ -1,6 +1,7 @@
 package flag
 
 import (
+	"github.com/docker/go-units"
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"golang.org/x/xerrors"
 
@@ -58,6 +59,12 @@ var (
 		Values:     xstrings.ToStringSlice(ftypes.AllImageSources),
 		Usage:      "image source(s) to use, in priority order",
 	}
+	MaxImageSize = Flag[string]{
+		Name:       "max-image-size",
+		ConfigName: "image.max-size",
+		Default:    "",
+		Usage:      "[EXPERIMENTAL] maximum image size to process, specified in a human-readable format (e.g., '44kB', '17MB'); an error will be returned if the image exceeds this size",
+	}
 )
 
 type ImageFlagGroup struct {
@@ -68,6 +75,7 @@ type ImageFlagGroup struct {
 	DockerHost          *Flag[string]
 	PodmanHost          *Flag[string]
 	ImageSources        *Flag[[]string]
+	MaxImageSize        *Flag[string]
 }
 
 type ImageOptions struct {
@@ -78,6 +86,7 @@ type ImageOptions struct {
 	DockerHost          string
 	PodmanHost          string
 	ImageSources        ftypes.ImageSources
+	MaxImageSize        int64
 }
 
 func NewImageFlagGroup() *ImageFlagGroup {
@@ -89,6 +98,7 @@ func NewImageFlagGroup() *ImageFlagGroup {
 		DockerHost:          DockerHostFlag.Clone(),
 		PodmanHost:          PodmanHostFlag.Clone(),
 		ImageSources:        SourceFlag.Clone(),
+		MaxImageSize:        MaxImageSize.Clone(),
 	}
 }
 
@@ -105,6 +115,7 @@ func (f *ImageFlagGroup) Flags() []Flagger {
 		f.DockerHost,
 		f.PodmanHost,
 		f.ImageSources,
+		f.MaxImageSize,
 	}
 }
 
@@ -124,6 +135,14 @@ func (f *ImageFlagGroup) ToOptions() (ImageOptions, error) {
 		}
 		platform = ftypes.Platform{Platform: pl}
 	}
+	var maxSize int64
+	if value := f.MaxImageSize.Value(); value != "" {
+		parsedSize, err := units.FromHumanSize(value)
+		if err != nil {
+			return ImageOptions{}, xerrors.Errorf("invalid max image size %q: %w", value, err)
+		}
+		maxSize = parsedSize
+	}
 
 	return ImageOptions{
 		Input:               f.Input.Value(),
@@ -133,5 +152,6 @@ func (f *ImageFlagGroup) ToOptions() (ImageOptions, error) {
 		DockerHost:          f.DockerHost.Value(),
 		PodmanHost:          f.PodmanHost.Value(),
 		ImageSources:        xstrings.ToTSlice[ftypes.ImageSource](f.ImageSources.Value()),
+		MaxImageSize:        maxSize,
 	}, nil
 }

pkg/flag/image_flags_test.go

@@ -0,0 +1,91 @@
+package flag_test
+
+import (
+	"testing"
+
+	"github.com/docker/go-units"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/aquasecurity/trivy/pkg/fanal/types"
+	"github.com/aquasecurity/trivy/pkg/flag"
+)
+
+func TestImageFlagGroup_ToOptions(t *testing.T) {
+	type fields struct {
+		maxImgSize string
+		platform   string
+	}
+	tests := []struct {
+		name    string
+		fields  fields
+		want    flag.ImageOptions
+		wantErr string
+	}{
+		{
+			name:   "happy default (without flags)",
+			fields: fields{},
+			want:   flag.ImageOptions{},
+		},
+		{
+			name: "happy path with max image size",
+			fields: fields{
+				maxImgSize: "10mb",
+			},
+			want: flag.ImageOptions{
+				MaxImageSize: units.MB * 10,
+			},
+		},
+		{
+			name: "invalid max image size",
+			fields: fields{
+				maxImgSize: "10foo",
+			},
+			wantErr: "invalid max image size",
+		},
+		{
+			name: "happy path with platform",
+			fields: fields{
+				platform: "linux/amd64",
+			},
+			want: flag.ImageOptions{
+				Platform: types.Platform{
+					Platform: &v1.Platform{
+						OS:           "linux",
+						Architecture: "amd64",
+					},
+				},
+			},
+		},
+		{
+			name: "invalid platform",
+			fields: fields{
+				platform: "unknown/unknown/unknown/unknown",
+			},
+			wantErr: "unable to parse platform",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Cleanup(viper.Reset)
+
+			setValue(flag.MaxImageSize.ConfigName, tt.fields.maxImgSize)
+			setValue(flag.PlatformFlag.ConfigName, tt.fields.platform)
+
+			f := &flag.ImageFlagGroup{
+				MaxImageSize: flag.MaxImageSize.Clone(),
+				Platform:     flag.PlatformFlag.Clone(),
+			}
+
+			got, err := f.ToOptions()
+			if tt.wantErr != "" {
+				assert.ErrorContains(t, err, tt.wantErr)
+				return
+			}
+			require.NoError(t, err)
+			assert.EqualExportedValues(t, tt.want, got)
+		})
+	}
+}

pkg/types/error.go

@@ -11,3 +11,12 @@ type ExitError struct {
 func (e *ExitError) Error() string {
 	return fmt.Sprintf("exit status %d", e.Code)
 }
+
+// UserError represents an error with a user-friendly message.
+type UserError struct {
+	Message string
+}
+
+func (e *UserError) Error() string {
+	return e.Message
+}