diff --git a/pkg/fanal/artifact/image/image.go b/pkg/fanal/artifact/image/image.go index cb991731d33b..f1b1a12f91cc 100644 --- a/pkg/fanal/artifact/image/image.go +++ b/pkg/fanal/artifact/image/image.go @@ -15,6 +15,7 @@ import ( "github.com/cheggaaa/pb/v3" "github.com/docker/go-units" v1 "github.com/google/go-containerregistry/pkg/v1" + "github.com/google/go-containerregistry/pkg/v1/partial" "github.com/samber/lo" "golang.org/x/xerrors" @@ -222,7 +223,7 @@ func (a Artifact) checkImageSize(ctx context.Context) error { return nil } - imageSize, err := a.imageSize(ctx) + imageSize, err := a.downloadImage(ctx) if err != nil { return xerrors.Errorf("failed to calculate image size: %w", err) } @@ -239,14 +240,49 @@ func (a Artifact) checkImageSize(ctx context.Context) error { return nil } -func (a Artifact) imageSize(ctx context.Context) (int64, error) { +// progressLayer wraps a v1.Layer to add progress bar functionality +type progressLayer struct { + v1.Layer + bar *pb.ProgressBar +} + +func newProgressLayer(layer v1.Layer, bar *pb.ProgressBar) (v1.Layer, error) { + return partial.CompressedToLayer(&progressLayer{ + Layer: layer, bar: bar, + }) +} + +func (l *progressLayer) Compressed() (io.ReadCloser, error) { + rc, err := l.Layer.Compressed() + if err != nil { + return nil, err + } + + return l.bar.NewProxyReader(rc), nil +} + +func (a Artifact) downloadImage(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 + if !a.artifactOption.NoProgress { + progressPool := pb.NewPool() + wrappedLayers, err := a.wrapLayers(layers, progressPool) + if err != nil { + return -1, xerrors.Errorf("failed to wrap") + } + if err := progressPool.Start(); err != nil { + log.Error("Failed to start progress bar pool", log.Err(err)) + } else { + defer progressPool.Stop() + layers = wrappedLayers + } + + } + var imageSize int64 p := parallel.NewPipeline(a.artifactOption.Parallel, false, layers, func(ctx context.Context, layer v1.Layer) (int64, error) { layerSize, err := a.downloadLayer(layer) @@ -268,6 +304,24 @@ func (a Artifact) imageSize(ctx context.Context) (int64, error) { return imageSize, nil } +func (a Artifact) wrapLayers(layers []v1.Layer, progressPool *pb.Pool) ([]v1.Layer, error) { + wrappedLayers := make([]v1.Layer, 0, len(layers)) + for _, l := range layers { + size, err := l.Size() + if err != nil { + return nil, err + } + bar := pb.New64(size).SetTemplate(pb.Full) + progressPool.Add(bar) + pl, err := newProgressLayer(l, bar) + if err != nil { + return nil, xerrors.Errorf("failed to create progress layer: %w", err) + } + wrappedLayers = append(wrappedLayers, pl) + } + return wrappedLayers, nil +} + func (a Artifact) downloadLayer(layer v1.Layer) (int64, error) { rc, err := layer.Compressed() if err != nil { @@ -287,21 +341,7 @@ func (a Artifact) downloadLayer(layer v1.Layer) (int64, error) { } defer f.Close() - 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) + dr, err := uncompressed(rc) if err != nil { return -1, xerrors.Errorf("failed to init decompressor: %w", err) } diff --git a/pkg/fanal/image/daemon/image.go b/pkg/fanal/image/daemon/image.go index d650851ad242..de680cf27312 100644 --- a/pkg/fanal/image/daemon/image.go +++ b/pkg/fanal/image/daemon/image.go @@ -151,6 +151,13 @@ func (img *image) configFile() (*v1.ConfigFile, error) { return img.Image.ConfigFile() } +func (img *image) Layers() ([]v1.Layer, error) { + if err := img.populateImage(); err != nil { + return nil, xerrors.Errorf("unable to populate: %w", err) + } + return img.Image.Layers() +} + func (img *image) LayerByDiffID(h v1.Hash) (v1.Layer, error) { if err := img.populateImage(); err != nil { return nil, xerrors.Errorf("unable to populate: %w", err) @@ -173,6 +180,20 @@ func (img *image) RepoDigests() []string { return img.inspect.RepoDigests } +// func (img *image) Digest() (v1.Hash, error) { +// return v1.Hash{}, nil +// } + +// func (img *image) LayerByDigest(v1.Hash) (v1.Layer, error) { +// return nil, nil +// } + +// func (img *image) Manifest() (*v1.Manifest, error) { +// return nil, nil +// } + +// func (img *image) MediaType() types.Med + func (img *image) diffIDs() ([]v1.Hash, error) { var diffIDs []v1.Hash for _, l := range img.inspect.RootFS.Layers {