🧹 split loading container images from opening connections

Signed-off-by: Ivan Milchev <[email protected]>

providers/os/connection/container/image/docker.go

@@ -67,20 +67,27 @@ func writeCompressedTarImage(img v1.Image, digest string) (*os.File, error) {
 	}
 	filename := f.Name()
 
-	ref, err := name.ParseReference(digest, name.WeakValidation)
-	if err != nil {
+	if err := WriteCompressedTarImageToFile(img, digest, f); err != nil {
 		os.Remove(filename)
 		return nil, err
+
+	}
+	return f, nil
+}
+
+func WriteCompressedTarImageToFile(img v1.Image, digest string, f *os.File) error {
+	ref, err := name.ParseReference(digest, name.WeakValidation)
+	if err != nil {
+		return err
 	}
 
 	err = tarball.Write(ref, img, f)
 	if err != nil {
-		os.Remove(filename)
-		return nil, err
+		return err
 	}
 
 	// Rewind, to later read the complete file for uncompress
 	f.Seek(0, io.SeekStart)
 
-	return f, nil
+	return nil
 }

providers/os/connection/container/image/registry.go

@@ -6,7 +6,6 @@ package image
 import (
 	"crypto/tls"
 	"fmt"
-	"io"
 	"net"
 	"net/http"
 	"strings"
@@ -74,14 +73,14 @@ func GetImageDescriptor(ref name.Reference, opts ...Option) (*remote.Descriptor,
 	return remote.Get(ref, remote.WithAuth(o.auth))
 }
 
-func LoadImageFromRegistry(ref name.Reference, opts ...Option) (v1.Image, io.ReadCloser, error) {
+func LoadImageFromRegistry(ref name.Reference, opts ...Option) (v1.Image, error) {
 	o := &options{
 		insecure: false,
 	}
 
 	for _, option := range opts {
 		if err := option(o); err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 	}
 
@@ -98,7 +97,7 @@ func LoadImageFromRegistry(ref name.Reference, opts ...Option) (v1.Image, io.Rea
 		auth, err := kc.Resolve(ref.Context())
 		if err != nil {
 			fmt.Printf("getting creds for %q: %v", ref, err)
-			return nil, nil, err
+			return nil, err
 		}
 		o.auth = auth
 	}
@@ -126,15 +125,7 @@ func LoadImageFromRegistry(ref name.Reference, opts ...Option) (v1.Image, io.Rea
 
 	img, err := remote.Image(ref, remote.WithAuth(o.auth), remote.WithTransport(tr))
 	if err != nil {
-		return nil, nil, err
-	}
-
-	// write image to disk (conmpressed, unflattened)
-	// Otherwise we can not later recognize it as a valid image
-	f, err := writeCompressedTarImage(img, ref.String())
-	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
-
-	return img, f, nil
+	return img, nil
 }

providers/os/connection/docker_container.go

@@ -7,7 +7,6 @@ import (
 	"context"
 	"errors"
 	"io"
-	"os"
 	"strconv"
 	"strings"
 
@@ -186,30 +185,38 @@ func NewContainerRegistryImage(id uint32, conf *inventory.Config, asset *invento
 			registryOpts = append(registryOpts, remoteOpts[i])
 		}
 
+		var conn *TarConnection
 		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 {
+				var rc io.ReadCloser
 				// read image from disk
 				img, rc, err = image.LoadImageFromDisk(asset.Connections[0].Options[COMPRESSED_IMAGE])
 				if err != nil {
 					return nil, err
 				}
+
+				conn, err = NewWithReader(id, conf, asset, rc)
+				if err != nil {
+					return nil, err
+				}
 				loadedImage = true
 			}
 		}
 		if !loadedImage {
-			img, rc, err = image.LoadImageFromRegistry(ref, registryOpts...)
+			img, 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
+
+			conn, err = NewTarConnectionForContainer(id, conf, asset, img)
+			if err != nil {
+				return nil, err
+			}
 		}
 
 		var identifier string
@@ -218,10 +225,6 @@ func NewContainerRegistryImage(id uint32, conf *inventory.Config, asset *invento
 			identifier = containerid.MondooContainerImageID(hash.String())
 		}
 
-		conn, err := NewWithReader(id, conf, asset, rc)
-		if err != nil {
-			return nil, err
-		}
 		conn.PlatformIdentifier = identifier
 		conn.Metadata.Name = containerid.ShortContainerImageID(hash.String())
 

providers/os/connection/tar.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"io"
 	"os"
+	"sync"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"github.com/google/go-containerregistry/pkg/v1/mutate"
@@ -32,9 +33,11 @@ const (
 var _ shared.Connection = (*TarConnection)(nil)
 
 type TarConnection struct {
-	id    uint32
-	asset *inventory.Asset
-	conf  *inventory.Config
+	id        uint32
+	asset     *inventory.Asset
+	conf      *inventory.Config
+	fetchFn   func() (string, error)
+	fetchOnce sync.Once
 
 	Fs      *provider_tar.FS
 	CloseFN func()
@@ -83,7 +86,24 @@ func (p *TarConnection) RunCommand(command string) (*shared.Command, error) {
 	return &res, nil
 }
 
+func (p *TarConnection) EnsureLoaded() {
+	if p.fetchFn != nil {
+		p.fetchOnce.Do(func() {
+			f, err := p.fetchFn()
+			if err != nil {
+				log.Error().Err(err).Msg("tar> could not fetch tar file")
+				return
+			}
+			if err := p.LoadFile(f); err != nil {
+				log.Error().Err(err).Msg("tar> could not load tar file")
+				return
+			}
+		})
+	}
+}
+
 func (p *TarConnection) FileSystem() afero.Fs {
+	p.EnsureLoaded()
 	return p.Fs
 }
 
@@ -155,10 +175,43 @@ func (c *TarConnection) Runtime() string {
 	return c.PlatformRuntime
 }
 
+func NewTarConnectionForContainer(id uint32, conf *inventory.Config, asset *inventory.Asset, img v1.Image) (*TarConnection, error) {
+	f, err := cache.RandomFile()
+	if err != nil {
+		return nil, err
+	}
+
+	return &TarConnection{
+		id:    id,
+		asset: asset,
+		Fs:    provider_tar.NewFs(""),
+		fetchFn: func() (string, error) {
+			err = cache.StreamToTmpFile(mutate.Extract(img), f)
+			if err != nil {
+				os.Remove(f.Name())
+				return "", err
+			}
+			log.Warn().Msg("tar> extracted image to temporary file")
+			f.Seek(0, io.SeekStart)
+			asset.Connections[0].Options[FLATTENED_IMAGE] = f.Name()
+			return f.Name(), nil
+		},
+		CloseFN: func() {
+			log.Warn().Str("tar", f.Name()).Msg("tar> remove temporary tar file on connection close")
+			os.Remove(f.Name())
+		},
+		PlatformKind:    conf.Type,
+		PlatformRuntime: conf.Runtime,
+		conf:            conf,
+	}, nil
+}
+
+// TODO: this one is used by plain tar connection
 func NewTarConnection(id uint32, conf *inventory.Config, asset *inventory.Asset) (*TarConnection, error) {
 	return NewWithClose(id, conf, asset, nil)
 }
 
+// Used with docker snapshots
 // NewWithReader provides a tar provider from a container image stream
 func NewWithReader(id uint32, conf *inventory.Config, asset *inventory.Asset, rc io.ReadCloser) (*TarConnection, error) {
 	filename := ""

providers/os/provider/provider.go

@@ -289,7 +289,7 @@ func (s *Service) Shutdown(req *plugin.ShutdownReq) (*plugin.ShutdownRes, error)
 	for i := range s.runtimes {
 		runtime := s.runtimes[i]
 		if x, ok := runtime.Connection.(*connection.TarConnection); ok {
-			x.CloseFN()
+			x.Close()
 		}
 	}
 	return &plugin.ShutdownRes{}, nil