feat(image): add progress bar for image layer pulling

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

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

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
+}

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,

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()
+}