diff --git a/go.mod b/go.mod index e0e4b58566bf..a868fde8e669 100644 --- a/go.mod +++ b/go.mod @@ -68,6 +68,7 @@ require ( github.com/hashicorp/hcl/v2 v2.23.0 github.com/hashicorp/terraform-exec v0.21.0 github.com/in-toto/in-toto-golang v0.9.0 + github.com/klauspost/compress v1.17.11 github.com/knqyf263/go-apk-version v0.0.0-20200609155635-041fdbb8563f github.com/knqyf263/go-deb-version v0.0.0-20241115132648-6f4aee6ccd23 github.com/knqyf263/go-rpm-version v0.0.0-20220614171824-631e686d1075 @@ -284,7 +285,6 @@ require ( github.com/josharian/intern v1.0.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/kevinburke/ssh_config v1.2.0 // indirect - github.com/klauspost/compress v1.17.11 // indirect github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect github.com/letsencrypt/boulder v0.0.0-20240620165639-de9c06129bec // indirect diff --git a/pkg/fanal/artifact/image/compression.go b/pkg/fanal/artifact/image/compression.go new file mode 100644 index 000000000000..f9e7db703ff6 --- /dev/null +++ b/pkg/fanal/artifact/image/compression.go @@ -0,0 +1,85 @@ +package image + +import ( + "bufio" + "bytes" + "compress/gzip" + "errors" + "io" + + "github.com/klauspost/compress/zstd" + "golang.org/x/xerrors" + + xio "github.com/aquasecurity/trivy/pkg/x/io" +) + +// https://en.wikipedia.org/wiki/List_of_file_signatures +var ( + gzipMagicNumber = []byte{'\x1f', '\x8b'} + zstdMagicNumber = []byte{'\x28', '\xb5', '\x2f', '\xfd'} +) + +type decompressor struct { + magicNumber []byte + wrap func(io.Reader) (io.ReadCloser, error) +} + +var decompressors = []decompressor{ + { + magicNumber: gzipMagicNumber, + wrap: func(r io.Reader) (io.ReadCloser, error) { + gr, err := gzip.NewReader(r) + if err != nil { + return nil, xerrors.Errorf("failed to create gzip reader: %w", err) + } + return gr, nil + }, + }, + { + magicNumber: zstdMagicNumber, + wrap: func(r io.Reader) (io.ReadCloser, error) { + zr, err := zstd.NewReader(r) + if err != nil { + return nil, xerrors.Errorf("failed to create zstd reader: %w", err) + } + return zr.IOReadCloser(), nil + }, + }, +} + +// uncompressed checks if the reader contains compressed data and returns the decompressed reader +// or the original reader if the data is not compressed. +func uncompressed(rc io.Reader) (io.ReadCloser, error) { + br := bufio.NewReader(rc) + for _, dec := range decompressors { + ok, err := hasMagicNumber(br, dec.magicNumber) + if err != nil { + return nil, xerrors.Errorf("failed to check file header: %w", err) + } + + if ok { + return dec.wrap(br) + } + } + + // decompression not required + return &xio.ReadCloser{ + Reader: rc, + CloseFunc: func() error { return nil }, + }, nil +} + +type peekReader interface { + io.Reader + Peek(n int) ([]byte, error) +} + +func hasMagicNumber(pr peekReader, magicNumber []byte) (bool, error) { + b, err := pr.Peek(len(magicNumber)) + if errors.Is(err, io.EOF) { + return false, nil + } else if err != nil { + return false, err + } + return bytes.Equal(b, magicNumber), nil +} diff --git a/pkg/fanal/artifact/image/image.go b/pkg/fanal/artifact/image/image.go index a5d92d905e3f..cb991731d33b 100644 --- a/pkg/fanal/artifact/image/image.go +++ b/pkg/fanal/artifact/image/image.go @@ -12,6 +12,7 @@ import ( "strings" "sync" + "github.com/cheggaaa/pb/v3" "github.com/docker/go-units" v1 "github.com/google/go-containerregistry/pkg/v1" "github.com/samber/lo" @@ -101,7 +102,7 @@ func (a Artifact) Inspect(ctx context.Context) (artifact.Reference, error) { a.logger.Debug("Detected diff ID", log.Any("diff_ids", diffIDs)) defer os.RemoveAll(a.cacheDir) - if err := a.checkImageSize(ctx, diffIDs); err != nil { + if err := a.checkImageSize(ctx); err != nil { return artifact.Reference{}, err } @@ -215,13 +216,13 @@ func (a Artifact) consolidateCreatedBy(diffIDs, layerKeys []string, configFile * return layerKeyMap } -func (a Artifact) checkImageSize(ctx context.Context, diffIDs []string) error { +func (a Artifact) checkImageSize(ctx context.Context) error { maxSize := a.artifactOption.ImageOption.MaxImageSize if maxSize == 0 { return nil } - imageSize, err := a.imageSize(ctx, diffIDs) + imageSize, err := a.imageSize(ctx) if err != nil { return xerrors.Errorf("failed to calculate image size: %w", err) } @@ -238,12 +239,17 @@ func (a Artifact) checkImageSize(ctx context.Context, diffIDs []string) error { return nil } -func (a Artifact) imageSize(ctx context.Context, diffIDs []string) (int64, error) { +func (a Artifact) imageSize(ctx context.Context) (int64, error) { + layers, err := a.image.Layers() + if err != nil { + return -1, xerrors.Errorf("failed to get image layers: %w", err) + } + var imageSize int64 - p := parallel.NewPipeline(a.artifactOption.Parallel, false, diffIDs, - func(_ context.Context, diffID string) (int64, error) { - layerSize, err := a.saveLayer(diffID) + p := parallel.NewPipeline(a.artifactOption.Parallel, false, layers, + func(ctx context.Context, layer v1.Layer) (int64, error) { + layerSize, err := a.downloadLayer(layer) if err != nil { return -1, xerrors.Errorf("failed to save layer: %w", err) } @@ -262,20 +268,58 @@ func (a Artifact) imageSize(ctx context.Context, diffIDs []string) (int64, error return imageSize, nil } -func (a Artifact) saveLayer(diffID string) (int64, error) { - _, rc, err := a.uncompressedLayer(diffID) +func (a Artifact) downloadLayer(layer v1.Layer) (int64, error) { + rc, err := layer.Compressed() if err != nil { - return -1, xerrors.Errorf("unable to get uncompressed layer %s: %w", diffID, err) + return -1, xerrors.Errorf("failed to fetch the layer: %w", err) } defer rc.Close() - f, err := os.Create(filepath.Join(a.cacheDir, diffID)) + h, err := layer.DiffID() + if err != nil { + return -1, xerrors.Errorf("failed to get hash of layer: %w", err) + } + + file := filepath.Join(a.cacheDir, h.String()) + f, err := os.Create(file) if err != nil { return -1, xerrors.Errorf("failed to create a file: %w", err) } defer f.Close() - return io.Copy(f, rc) + size, err := layer.Size() + if err != nil { + return -1, xerrors.Errorf("size error: %w", err) + } + + bar := pb.Simple.Start64(size) + if a.artifactOption.NoProgress { + bar.SetWriter(io.Discard) + } + pr := bar.NewProxyReader(rc) + defer bar.Finish() + + bar.Set("prefix", shortenHash(h.Hex, 12)) + + dr, err := uncompressed(pr) + if err != nil { + return -1, xerrors.Errorf("failed to init decompressor: %w", err) + } + defer dr.Close() + + n, err := io.Copy(f, dr) + if err != nil { + return -1, xerrors.Errorf("failed to download layer: %w", err) + } + + return n, nil +} + +func shortenHash(hash string, length int) string { + if len(hash) > length { + return hash[:length] + } + return hash } func (a Artifact) inspect(ctx context.Context, missingImage string, layerKeys, baseDiffIDs []string, diff --git a/pkg/x/io/io.go b/pkg/x/io/io.go index 1fe3d6c0c2fc..bf4131d94e9c 100644 --- a/pkg/x/io/io.go +++ b/pkg/x/io/io.go @@ -71,3 +71,12 @@ type nopCloser struct { } func (nopCloser) Close() error { return nil } + +type ReadCloser struct { + io.Reader + CloseFunc func() error +} + +func (rac *ReadCloser) Close() error { + return rac.CloseFunc() +}