🐛 Fetch and save container images only once (#1938)

* 🐛 Fetch and save container images only once Fixes #1771 Signed-off-by: Christian Zunker <[email protected]> * Fix license header Signed-off-by: Christian Zunker <[email protected]> * Fix tests Signed-off-by: Christian Zunker <[email protected]> --------- Signed-off-by: Christian Zunker <[email protected]>
mondoohq · Sep 27, 2023 · 1418a74 · 1418a74
1 parent 5409771
commit 1418a74
Show file tree

Hide file tree

Showing 5 changed files with 97 additions and 26 deletions.
diff --git a/providers/os/connection/container/image/disk.go b/providers/os/connection/container/image/disk.go
@@ -0,0 +1,25 @@
+// Copyright (c) Mondoo, Inc.
+// SPDX-License-Identifier: BUSL-1.1
+
+package image
+
+import (
+	"io"
+	"os"
+
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/tarball"
+)
+
+func LoadImageFromDisk(filepath string) (v1.Image, io.ReadCloser, error) {
+	img, err := tarball.ImageFromPath(filepath, nil)
+	if err != nil {
+		return nil, nil, err
+	}
+	rc, err := os.Open(filepath)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return img, rc, nil
+}
diff --git a/providers/os/connection/container/image/docker.go b/providers/os/connection/container/image/docker.go
@@ -79,7 +79,7 @@ func writeCompressedTarImage(img v1.Image, digest string) (*os.File, error) {
 		return nil, err
 	}
 
-	// Rewindo, to later read the complete file for uncompress
+	// Rewind, to later read the complete file for uncompress
 	f.Seek(0, io.SeekStart)
 
 	return f, nil

diff --git a/providers/os/connection/docker_container.go b/providers/os/connection/docker_container.go
@@ -7,11 +7,13 @@ import (
 	"context"
 	"errors"
 	"io"
+	"os"
 	"strconv"
 	"strings"
 
 	"github.com/docker/docker/client"
 	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/afero"
 	"go.mondoo.com/cnquery/providers-sdk/v1/inventory"
@@ -184,9 +186,30 @@ func NewContainerRegistryImage(id uint32, conf *inventory.Config, asset *invento
 			registryOpts = append(registryOpts, remoteOpts[i])
 		}
 
-		img, rc, err := image.LoadImageFromRegistry(ref, registryOpts...)
-		if err != nil {
-			return nil, err
+		var img v1.Image
+		var rc io.ReadCloser
+		loadedImage := false
+		if asset.Connections[0].Options != nil {
+			if _, ok := asset.Connections[0].Options[COMPRESSED_IMAGE]; ok {
+				// read image from disk
+				img, rc, err = image.LoadImageFromDisk(asset.Connections[0].Options[COMPRESSED_IMAGE])
+				if err != nil {
+					return nil, err
+				}
+				loadedImage = true
+			}
+		}
+		if !loadedImage {
+			img, rc, err = image.LoadImageFromRegistry(ref, registryOpts...)
+			if err != nil {
+				return nil, err
+			}
+			if asset.Connections[0].Options == nil {
+				asset.Connections[0].Options = map[string]string{}
+			}
+			osFile := rc.(*os.File)
+			filename := osFile.Name()
+			asset.Connections[0].Options[COMPRESSED_IMAGE] = filename
 		}
 
 		var identifier string
@@ -195,7 +218,7 @@ func NewContainerRegistryImage(id uint32, conf *inventory.Config, asset *invento
 			identifier = containerid.MondooContainerImageID(hash.String())
 		}
 
-		conn, err := NewWithReader(id, conf, asset, rc, nil)
+		conn, err := NewWithReader(id, conf, asset, rc)
 		if err != nil {
 			return nil, err
 		}
@@ -340,7 +363,7 @@ func NewDockerContainerImageConnection(id uint32, conf *inventory.Config, asset
 	asset.Name = ii.Name
 	asset.Labels = ii.Labels
 
-	tarConn, err := NewWithReader(id, conf, asset, rc, nil)
+	tarConn, err := NewWithReader(id, conf, asset, rc)
 	if err != nil {
 		return nil, err
 	}

diff --git a/providers/os/connection/tar.go b/providers/os/connection/tar.go
@@ -24,8 +24,10 @@ import (
 )
 
 const (
-	Tar         shared.ConnectionType = "tar"
-	OPTION_FILE                       = "path"
+	Tar              shared.ConnectionType = "tar"
+	OPTION_FILE                            = "path"
+	FLATTENED_IMAGE                        = "flattened_path"
+	COMPRESSED_IMAGE                       = "compressed_path"
 )
 
 type TarConnection struct {
@@ -157,7 +159,7 @@ func NewTarConnection(id uint32, conf *inventory.Config, asset *inventory.Asset)
 }
 
 // NewWithReader provides a tar provider from a container image stream
-func NewWithReader(id uint32, conf *inventory.Config, asset *inventory.Asset, rc io.ReadCloser, close func()) (*TarConnection, error) {
+func NewWithReader(id uint32, conf *inventory.Config, asset *inventory.Asset, rc io.ReadCloser) (*TarConnection, error) {
 	filename := ""
 	if x, ok := rc.(*os.File); ok {
 		filename = x.Name()
@@ -185,7 +187,7 @@ func NewWithReader(id uint32, conf *inventory.Config, asset *inventory.Asset, rc
 			OPTION_FILE: filename,
 		},
 	}, asset, func() {
-		// remove temporary file on stream close
+		log.Debug().Str("tar", filename).Msg("tar> remove temporary tar file on connection close")
 		os.Remove(filename)
 	})
 }
@@ -208,16 +210,11 @@ func NewWithClose(id uint32, conf *inventory.Config, asset *inventory.Asset, clo
 		identifier = containerid.MondooContainerImageID(hash.String())
 
 		// we cache the flattened image locally
-		c, err := newWithFlattenedImage(id, conf, asset, &img)
+		c, err := newWithFlattenedImage(id, conf, asset, &img, closeFn)
 		if err != nil {
 			return nil, err
 		}
 
-		// remove unflattened image file, we now have a flattened image
-		if closeFn != nil {
-			closeFn()
-		}
-
 		c.PlatformIdentifier = identifier
 		return c, nil
 	} else {
@@ -247,24 +244,39 @@ func NewWithClose(id uint32, conf *inventory.Config, asset *inventory.Asset, clo
 	}
 }
 
-func newWithFlattenedImage(id uint32, conf *inventory.Config, asset *inventory.Asset, img *v1.Image) (*TarConnection, error) {
-	f, err := cache.RandomFile()
-	if err != nil {
-		return nil, err
+func newWithFlattenedImage(id uint32, conf *inventory.Config, asset *inventory.Asset, img *v1.Image, closeFn func()) (*TarConnection, error) {
+	imageFilename := ""
+	useCached := false
+	if asset != nil && len(asset.Connections) > 0 {
+		if x, ok := asset.Connections[0].Options[FLATTENED_IMAGE]; ok && x != "" {
+			log.Debug().Str("tar", asset.Connections[0].Options[FLATTENED_IMAGE]).Msg("tar> use cached tar file")
+			imageFilename = asset.Connections[0].Options[FLATTENED_IMAGE]
+			useCached = true
+		}
 	}
-	imageFilename := f.Name()
-	err = cache.StreamToTmpFile(mutate.Extract(*img), f)
-	if err != nil {
-		os.Remove(imageFilename)
-		return nil, err
+	if !useCached {
+		f, err := cache.RandomFile()
+		if err != nil {
+			return nil, err
+		}
+		imageFilename = f.Name()
+		err = cache.StreamToTmpFile(mutate.Extract(*img), f)
+		if err != nil {
+			os.Remove(imageFilename)
+			return nil, err
+		}
 	}
 
 	c := &TarConnection{
 		id:    id,
 		asset: asset,
 		Fs:    provider_tar.NewFs(imageFilename),
 		CloseFN: func() {
+			if closeFn != nil {
+				closeFn()
+			}
 			// remove temporary file on stream close
+			log.Debug().Str("tar", imageFilename).Msg("tar> remove temporary flattened image file on connection close")
 			os.Remove(imageFilename)
 		},
 		PlatformKind:    conf.Backend,
@@ -275,10 +287,17 @@ func newWithFlattenedImage(id uint32, conf *inventory.Config, asset *inventory.A
 			},
 		},
 	}
+	if asset != nil && len(asset.Connections) > 0 {
+		if asset.Connections[0].Options == nil {
+			asset.Connections[0].Options = map[string]string{}
+		}
+		asset.Connections[0].Options[FLATTENED_IMAGE] = imageFilename
+	}
 
-	err = c.LoadFile(imageFilename)
+	err := c.LoadFile(imageFilename)
 	if err != nil {
 		log.Error().Err(err).Str("tar", imageFilename).Msg("tar> could not load tar file")
+		os.Remove(imageFilename)
 		return nil, err
 	}
 

diff --git a/providers/os/resources/discovery/container_registry/resolver.go b/providers/os/resources/discovery/container_registry/resolver.go
@@ -61,6 +61,10 @@ func (r *Resolver) Resolve(ctx context.Context, root *inventory.Asset, conf *inv
 		if err != nil {
 			return nil, err
 		}
+		// keep already set options, i.e. image paths
+		if conf.Options != nil && a.Connections[0].Options == nil {
+			a.Connections[0].Options = conf.Options
+		}
 
 		if conf.Insecure {
 			for i := range a.Connections {