cri-o · sohankunkerkar · Apr 17, 2025 · Apr 23, 2025 · Apr 24, 2025 · Apr 25, 2025
@@ -43,6 +43,7 @@ require (
 	github.com/intel/goresctrl v0.8.0
 	github.com/json-iterator/go v1.1.12
 	github.com/kata-containers/kata-containers/src/runtime v0.0.0-20240208092920-b99f57452225
+	github.com/klauspost/compress v1.18.0
 	github.com/moby/sys/mountinfo v0.7.2
 	github.com/moby/sys/user v0.3.0
 	github.com/moby/sys/userns v0.1.0
@@ -162,7 +163,6 @@ require (
 	github.com/jinzhu/copier v0.4.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
-	github.com/klauspost/compress v1.18.0 // indirect
 	github.com/klauspost/pgzip v1.2.6 // indirect
 	github.com/knqyf263/go-plugin v0.8.1-0.20241122101054-d8d42059d8f1 // indirect
 	github.com/letsencrypt/boulder v0.0.0-20240620165639-de9c06129bec // indirect

@@ -77,6 +77,8 @@
 	runtimePath           string // runtime path for a given platform
 	execPIDs              map[int]bool
 	runtimeUser           *types.ContainerUser
+	cleanups              []func()
+	cleanupOnce           *sync.Once
 }
 
 func (c *Container) CRIAttributes() *types.ContainerAttributes {
@@ -877,3 +879,18 @@
 func (c *Container) RuntimeUser() *types.ContainerUser {
 	return c.runtimeUser
 }
+
+func (c *Container) AddCleanup(cleanup func()) {
+	c.cleanups = append(c.cleanups, cleanup)
+}
+
+func (c *Container) Cleanup() {
+	if c.cleanupOnce == nil {
+		c.cleanupOnce = &sync.Once{}
+	}
+	c.cleanupOnce.Do(func() {
+		for _, cleanup := range c.cleanups {
+			cleanup()
+		}
+	})
+}
@@ -919,6 +919,7 @@ func (r *runtimeOCI) StopLoopForContainer(c *Container, bm kwait.BackoffManager)
 		c.state.Finished = time.Now()
 		c.opLock.Unlock()
 		c.SetAsDoneStopping()
+		c.Cleanup()
 	}()
 
 	if c.state.Status == ContainerStatePaused {

@@ -689,6 +689,7 @@
 		return err
 	}
 
+	c.Cleanup()
 	c.state.Finished = time.Now()
 
 	return nil

@@ -1,14 +1,18 @@
 package ociartifact
 
 import (
+	"archive/tar"
+	"compress/gzip"
 	"context"
 	"encoding/hex"
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
 
 	modelSpec "github.com/CloudNativeAI/model-spec/specs-go/v1"
 	"github.com/containers/common/libimage"
@@ -17,6 +21,7 @@
 	"github.com/containers/image/v5/oci/layout"
 	"github.com/containers/image/v5/pkg/blobinfocache"
 	"github.com/containers/image/v5/types"
+	"github.com/klauspost/compress/zstd"
 	"github.com/opencontainers/go-digest"
 	specs "github.com/opencontainers/image-spec/specs-go/v1"
 
@@ -469,7 +474,16 @@
 	// Source path of the blob, i.e. full path in the blob dir.
 	SourcePath string
 	// Name of the file in the artifact.
-	Name string
+	Name        string
+	cleanupFunc func()
+	cleanupOnce *sync.Once
+}
+
+// Cleanup triggers the deletion of the extracted directory.
+func (b *BlobMountPath) Cleanup() {
+	if b.cleanupFunc != nil && b.cleanupOnce != nil {
+		b.cleanupOnce.Do(b.cleanupFunc)
+	}
 }
 
 // BlobMountPaths retrieves the local file paths for all blobs in the provided artifact and returns them as BlobMountPath slices.
@@ -479,35 +493,131 @@
 		return nil, fmt.Errorf("failed to get an image reference: %w", err)
 	}
 
+	extractDir := filepath.Join(s.rootPath, "extracted", artifact.Digest().Encoded())
+	if err := os.MkdirAll(extractDir, 0o755); err != nil {
+		return nil, fmt.Errorf("failed to create artifact dir: %w", err)
+	}
+
+	cleanupFunc := func() { os.RemoveAll(extractDir) }
+	cleanupOnce := &sync.Once{}
+
+	// Cleanup on function exit if we encounter an error
+	cleanup := true
+	defer func() {
+		if cleanup {
+			os.RemoveAll(extractDir)
+		}
+	}()
+
 	src, err := ref.NewImageSource(ctx, sys)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get an image source: %w", err)
 	}
-
 	defer src.Close()
 
 	mountPaths := make([]BlobMountPath, 0, len(artifact.Manifest().Layers))
+	bic := blobinfocache.DefaultCache(s.systemContext)
 
 	for _, l := range artifact.Manifest().Layers {
-		path, err := layout.GetLocalBlobPath(ctx, src, l.Digest)
+		err = func() error {
+			blob, _, err := s.impl.GetBlob(ctx, src, types.BlobInfo{Digest: l.Digest}, bic)
+			if err != nil {
+				return fmt.Errorf("failed to get blob: %w", err)
+			}
+			defer blob.Close()
+
+			content, err := s.processLayerContent(l.MediaType, blob)
+			if err != nil {
+				return fmt.Errorf("failed to process layer: %w", err)
+			}
+
+			name := artifactName(l.Annotations)
+			if name == "" {
+				log.Warnf(ctx, "Unable to find name for artifact layer which makes it not mountable")
+
+				return nil
+			}
+
+			filePath := filepath.Join(extractDir, name)
+			if err := os.WriteFile(filePath, content, 0o644); err != nil {
+				log.Errorf(ctx, "Failed to write file: %v", err)
+				os.Remove(filePath)
+				return err
+			}
+
+			mountPaths = append(mountPaths, BlobMountPath{
+				SourcePath:  filePath,
+				Name:        name,
+				cleanupFunc: cleanupFunc,
+				cleanupOnce: cleanupOnce,
+			})
+			return nil
+		}()
+		// If any layer extraction fails, abort the entire operation.
+		if err != nil {
+			return nil, fmt.Errorf("failed to extract artifact layer: %w", err)
+		}
+	}
+	// Disable cleanup on success.
+	// mountArtifact() will handle the cleanup.
+	cleanup = false
+	return mountPaths, nil
+}
+
+// processLayerContent decompresses and processes layer content based on media type.
+// Supports gzip, zstd, and uncompressed tar formats. Returns raw bytes for non-tar content.
+// Instead of relying on the exact media type like "application/vnd.oci.image.layer.v1.tar+zstd",
+// it would be helpful to stick with more generic patterns like ".tar+gzip" or ".tar+zstd".
+func (s *Store) processLayerContent(mediaType string, r io.Reader) ([]byte, error) {
+	switch {
+	case strings.HasSuffix(mediaType, ".tar+gzip"):
+		gr, err := gzip.NewReader(r)
 		if err != nil {
-			return nil, fmt.Errorf("failed to get a local blob path: %w", err)
+			return nil, fmt.Errorf("gzip decompress failed: %w", err)
 		}
+		defer gr.Close()
 
-		name := artifactName(l.Annotations)
-		if name == "" {
-			log.Warnf(ctx, "Unable to find name for artifact layer which makes it not mountable")
+		return s.extractSingleFileFromTar(gr)
 
-			continue
+	case strings.HasSuffix(mediaType, ".tar+zstd"):
+		zr, err := zstd.NewReader(r)
+		if err != nil {
+			return nil, fmt.Errorf("zstd decompress failed: %w", err)
 		}
+		defer zr.Close()
 
-		mountPaths = append(mountPaths, BlobMountPath{
-			SourcePath: path,
-			Name:       name,
-		})
+		return s.extractSingleFileFromTar(zr)
+
+	case strings.HasSuffix(mediaType, ".tar"):
+		return s.extractSingleFileFromTar(r)
+
+	default:
+		data, err := io.ReadAll(r)
+		if err != nil {
+			return nil, fmt.Errorf("reading blob content: %w", err)
+		}
+
+		return data, nil
 	}
+}
 
-	return mountPaths, nil
+func (s *Store) extractSingleFileFromTar(r io.Reader) ([]byte, error) {
+	tr := tar.NewReader(r)
+
+	for {
+		hdr, err := tr.Next()
+		if err == io.EOF {
+			return nil, errors.New("no files found in tar archive")
+		}
+
+		if err != nil {
+			return nil, fmt.Errorf("reading tar header: %w", err)
+		}
+
+		if hdr.Typeflag == tar.TypeReg {
+			return io.ReadAll(tr)
+		}
+	}
 }
 
 func artifactName(annotations map[string]string) string {

@@ -837,7 +837,7 @@
 		}
 	}()
 
-	containerVolumes, ociMounts, safeMounts, err := s.addOCIBindMounts(ctx, ctr, mountLabel, s.config.BindMountPrefix, s.config.AbsentMountSourcesToReject, maybeRelabel, skipRelabel, cgroup2RW, idMapSupport, rroSupport, s.ContainerServer.Config().Root, containerInfo.RunDir)
+	containerVolumes, ociMounts, safeMounts, cleanups, err := s.addOCIBindMounts(ctx, ctr, mountLabel, s.config.BindMountPrefix, s.config.AbsentMountSourcesToReject, maybeRelabel, skipRelabel, cgroup2RW, idMapSupport, rroSupport, s.ContainerServer.Config().Root, containerInfo.RunDir)
 	if err != nil {
 		return nil, err
 	}
@@ -1380,6 +1380,10 @@
 		ociContainer.AddVolume(cv)
 	}
 
+	for _, cleanup := range cleanups {
+		ociContainer.AddCleanup(cleanup)
+	}
+
 	return ociContainer, nil
 }