fix: harden image manager and split CI tests by platform

Amp-Thread-ID: https://ampcode.com/threads/T-019c7d89-b268-73f7-9375-68d0c1c83103

Author: lox

.github/workflows/ci.yml

@@ -14,9 +14,11 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os:
-          - ubuntu-latest
-          - macos-latest
+        include:
+          - os: ubuntu-latest
+            test_task: test-full
+          - os: macos-latest
+            test_task: test
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -30,4 +32,4 @@ jobs:
         run: mise run build
 
       - name: Test
-        run: mise run test
+        run: mise run ${{ matrix.test_task }}

.mise.toml

@@ -5,7 +5,11 @@ go = "1.23"
 GOTOOLCHAIN = "local"
 
 [tasks.test]
-description = "Run Go test suite"
+description = "Run cross-platform Go test suite"
+run = "go test ./..."
+
+[tasks.test-full]
+description = "Run full Go test suite"
 run = "go test ./..."
 
 [tasks.build]

internal/backend/firecracker/backend.go

@@ -25,13 +25,28 @@ import (
 	fcvsock "github.com/firecracker-microvm/firecracker-go-sdk/vsock"
 )
 
-type Adapter struct{}
+type imageEnsurer interface {
+	Ensure(context.Context, string) (imagemgr.EnsureResult, error)
+}
+
+type imageManagerFactory func() (imageEnsurer, error)
+
+type Adapter struct {
+	imageManagerOnce sync.Once
+	imageManager     imageEnsurer
+	imageManagerErr  error
+	newImageManager  imageManagerFactory
+}
 
 const runObservabilityFile = "run-observability.json"
 const vsockDialRetryInterval = 50 * time.Millisecond
 
 func New() *Adapter {
-	return &Adapter{}
+	return &Adapter{newImageManager: defaultImageManagerFactory}
+}
+
+func defaultImageManagerFactory() (imageEnsurer, error) {
+	return imagemgr.New(imagemgr.Options{})
 }
 
 func (a *Adapter) Name() string {
@@ -280,7 +295,7 @@ func (a *Adapter) run(ctx context.Context, req backend.RunRequest, stream backen
 		return nil, fmt.Errorf("kernel image %s: %w", kernelPath, err)
 	}
 
-	imageArtifact, err := ensureImageArtifact(ctx, req.Policy.ImageRef)
+	imageArtifact, err := a.ensureImageArtifact(ctx, req.Policy.ImageRef)
 	if err != nil {
 		return nil, err
 	}
@@ -570,13 +585,13 @@ type imageArtifact struct {
 	CacheHit   bool
 }
 
-func ensureImageArtifact(ctx context.Context, imageRef string) (imageArtifact, error) {
+func (a *Adapter) ensureImageArtifact(ctx context.Context, imageRef string) (imageArtifact, error) {
 	trimmedRef := strings.TrimSpace(imageRef)
 	if trimmedRef == "" {
 		return imageArtifact{}, errors.New("sandbox.image.ref is required for launched execution")
 	}
 
-	mgr, err := imagemgr.New(imagemgr.Options{})
+	mgr, err := a.getImageManager()
 	if err != nil {
 		return imageArtifact{}, fmt.Errorf("initialise image manager: %w", err)
 	}
@@ -594,6 +609,22 @@ func ensureImageArtifact(ctx context.Context, imageRef string) (imageArtifact, e
 	}, nil
 }
 
+func (a *Adapter) getImageManager() (imageEnsurer, error) {
+	if a.newImageManager == nil {
+		a.newImageManager = defaultImageManagerFactory
+	}
+	a.imageManagerOnce.Do(func() {
+		a.imageManager, a.imageManagerErr = a.newImageManager()
+	})
+	if a.imageManagerErr != nil {
+		return nil, a.imageManagerErr
+	}
+	if a.imageManager == nil {
+		return nil, errors.New("image manager factory returned nil manager")
+	}
+	return a.imageManager, nil
+}
+
 func runResultMessage(base string) string {
 	return base + "; rootfs writes discarded after run"
 }

internal/imagemgr/imagemgr.go

@@ -328,18 +328,22 @@ func (m *Manager) persistFromTarStream(ctx context.Context, req persistFromTarRe
 	}
 
 	outputPath := filepath.Join(m.cacheDir, strings.TrimPrefix(req.Digest, "sha256:")+".ext4")
-	tmpPath := outputPath + ".tmp"
-	if err := os.Remove(tmpPath); err != nil && !os.IsNotExist(err) {
-		return Record{}, fmt.Errorf("remove stale image temp file %q: %w", tmpPath, err)
+	tmpFile, err := os.CreateTemp(m.cacheDir, strings.TrimPrefix(req.Digest, "sha256:")+".tmp-*.ext4")
+	if err != nil {
+		return Record{}, fmt.Errorf("create temporary image artifact for %q: %w", req.Digest, err)
+	}
+	tmpPath := tmpFile.Name()
+	if err := tmpFile.Close(); err != nil {
+		_ = os.Remove(tmpPath)
+		return Record{}, fmt.Errorf("close temporary image artifact file %q: %w", tmpPath, err)
 	}
+	defer os.Remove(tmpPath)
 
 	sizeBytes, err := m.materialize(ctx, req.TarStream, tmpPath)
 	if err != nil {
-		_ = os.Remove(tmpPath)
 		return Record{}, err
 	}
 	if err := os.Rename(tmpPath, outputPath); err != nil {
-		_ = os.Remove(tmpPath)
 		return Record{}, fmt.Errorf("move image artifact to cache %q: %w", outputPath, err)
 	}
 
@@ -354,6 +358,7 @@ func (m *Manager) persistFromTarStream(ctx context.Context, req persistFromTarRe
 		OCIConfig:  req.OCIConfig,
 	}
 	if err := m.upsertRecord(ctx, record); err != nil {
+		_ = os.Remove(outputPath)
 		return Record{}, err
 	}
 

internal/imagemgr/imagemgr_test.go

@@ -4,6 +4,7 @@ import (
 	"archive/tar"
 	"bytes"
 	"context"
+	"errors"
 	"io"
 	"os"
 	"path/filepath"
@@ -119,6 +120,110 @@ func TestRemoveByRefSelector(t *testing.T) {
 	}
 }
 
+func TestPersistFromTarStreamUsesUniqueTempPaths(t *testing.T) {
+	t.Parallel()
+
+	cacheDir := filepath.Join(t.TempDir(), "cache")
+	dbPath := filepath.Join(t.TempDir(), "state", "metadata.db")
+	if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
+		t.Fatalf("create state dir: %v", err)
+	}
+
+	var outputPaths []string
+	manager, err := New(Options{
+		CacheDir:       cacheDir,
+		MetadataDBPath: dbPath,
+		MaterializeRootFS: func(_ context.Context, _ io.Reader, outputPath string) (int64, error) {
+			outputPaths = append(outputPaths, outputPath)
+			return 0, errors.New("materialise fail")
+		},
+	})
+	if err != nil {
+		t.Fatalf("create manager: %v", err)
+	}
+
+	now := time.Unix(1_700_000_001, 0).UTC()
+	req := persistFromTarRequest{
+		Ref:        testImageRef,
+		Digest:     "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+		TarStream:  bytes.NewReader(testRootFSTar(t)),
+		Source:     "import",
+		CreatedAt:  now,
+		LastUsedAt: now,
+	}
+
+	if _, err := manager.persistFromTarStream(context.Background(), req); err == nil {
+		t.Fatal("expected first persistFromTarStream to fail")
+	}
+	req.TarStream = bytes.NewReader(testRootFSTar(t))
+	if _, err := manager.persistFromTarStream(context.Background(), req); err == nil {
+		t.Fatal("expected second persistFromTarStream to fail")
+	}
+
+	if len(outputPaths) != 2 {
+		t.Fatalf("expected two materialise attempts, got %d", len(outputPaths))
+	}
+	if outputPaths[0] == outputPaths[1] {
+		t.Fatalf("expected unique temporary output paths, got %q twice", outputPaths[0])
+	}
+}
+
+func TestPersistFromTarStreamRemovesArtifactWhenMetadataWriteFails(t *testing.T) {
+	t.Parallel()
+
+	cacheDir := filepath.Join(t.TempDir(), "cache")
+	stateDir := filepath.Join(t.TempDir(), "state")
+	dbPath := filepath.Join(stateDir, "metadata.db")
+	if err := os.MkdirAll(stateDir, 0o755); err != nil {
+		t.Fatalf("create state dir: %v", err)
+	}
+
+	manager, err := New(Options{
+		CacheDir:       cacheDir,
+		MetadataDBPath: dbPath,
+		MaterializeRootFS: func(_ context.Context, _ io.Reader, outputPath string) (int64, error) {
+			if err := os.WriteFile(outputPath, []byte("fake-ext4"), 0o644); err != nil {
+				return 0, err
+			}
+			if err := os.Chmod(dbPath, 0o444); err != nil {
+				return 0, err
+			}
+			if err := os.Chmod(stateDir, 0o555); err != nil {
+				return 0, err
+			}
+			return int64(len("fake-ext4")), nil
+		},
+	})
+	if err != nil {
+		t.Fatalf("create manager: %v", err)
+	}
+	t.Cleanup(func() {
+		_ = os.Chmod(stateDir, 0o755)
+		_ = os.Chmod(dbPath, 0o644)
+	})
+
+	digest := "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
+	now := time.Unix(1_700_000_002, 0).UTC()
+	req := persistFromTarRequest{
+		Ref:        "ghcr.io/buildkite/cleanroom-base/alpine@" + digest,
+		Digest:     digest,
+		TarStream:  bytes.NewReader(testRootFSTar(t)),
+		Source:     "import",
+		CreatedAt:  now,
+		LastUsedAt: now,
+	}
+
+	_, err = manager.persistFromTarStream(context.Background(), req)
+	if err == nil {
+		t.Fatal("expected persistFromTarStream to fail when metadata write fails")
+	}
+
+	artifactPath := filepath.Join(cacheDir, strings.TrimPrefix(digest, "sha256:")+".ext4")
+	if _, statErr := os.Stat(artifactPath); !os.IsNotExist(statErr) {
+		t.Fatalf("expected artifact %s to be removed after metadata failure, stat err=%v", artifactPath, statErr)
+	}
+}
+
 func newTestManager(t *testing.T, pullFn func(context.Context, string) (io.ReadCloser, OCIConfig, error)) *Manager {
 	t.Helper()
 

internal/imagemgr/materialize.go

@@ -3,6 +3,7 @@ package imagemgr
 import (
 	"archive/tar"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -125,13 +126,22 @@ func extractTar(root string, stream io.Reader) error {
 
 		switch hdr.Typeflag {
 		case tar.TypeDir:
+			if err := ensureNoSymlinkPath(root, targetPath, true); err != nil {
+				return err
+			}
 			if err := os.MkdirAll(targetPath, os.FileMode(hdr.Mode)); err != nil {
 				return fmt.Errorf("create directory %q from tar stream: %w", targetPath, err)
 			}
 		case tar.TypeReg, tar.TypeRegA:
+			if err := ensureNoSymlinkPath(root, filepath.Dir(targetPath), true); err != nil {
+				return err
+			}
 			if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
 				return fmt.Errorf("create parent directory for %q: %w", targetPath, err)
 			}
+			if err := ensureNoSymlinkPath(root, targetPath, true); err != nil {
+				return err
+			}
 			f, err := os.OpenFile(targetPath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.FileMode(hdr.Mode))
 			if err != nil {
 				return fmt.Errorf("create file %q from tar stream: %w", targetPath, err)
@@ -144,9 +154,18 @@ func extractTar(root string, stream io.Reader) error {
 				return fmt.Errorf("close file %q from tar stream: %w", targetPath, err)
 			}
 		case tar.TypeSymlink:
+			if err := ensureNoSymlinkPath(root, filepath.Dir(targetPath), true); err != nil {
+				return err
+			}
 			if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
 				return fmt.Errorf("create parent directory for symlink %q: %w", targetPath, err)
 			}
+			if err := validateSymlinkTarget(root, targetPath, hdr.Linkname); err != nil {
+				return err
+			}
+			if err := os.Remove(targetPath); err != nil && !os.IsNotExist(err) {
+				return fmt.Errorf("remove existing symlink path %q: %w", targetPath, err)
+			}
 			if err := os.Symlink(hdr.Linkname, targetPath); err != nil && !os.IsExist(err) {
 				return fmt.Errorf("create symlink %q -> %q from tar stream: %w", targetPath, hdr.Linkname, err)
 			}
@@ -155,9 +174,21 @@ func extractTar(root string, stream io.Reader) error {
 			if err != nil {
 				return err
 			}
+			if err := ensureNoSymlinkPath(root, linkTarget, false); err != nil {
+				return err
+			}
+			if err := ensureNoSymlinkPath(root, filepath.Dir(targetPath), true); err != nil {
+				return err
+			}
 			if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
 				return fmt.Errorf("create parent directory for hard link %q: %w", targetPath, err)
 			}
+			if err := ensureNoSymlinkPath(root, targetPath, true); err != nil {
+				return err
+			}
+			if err := os.Remove(targetPath); err != nil && !os.IsNotExist(err) {
+				return fmt.Errorf("remove existing hardlink path %q: %w", targetPath, err)
+			}
 			if err := os.Link(linkTarget, targetPath); err != nil {
 				return fmt.Errorf("create hard link %q -> %q from tar stream: %w", targetPath, linkTarget, err)
 			}
@@ -182,3 +213,54 @@ func safeJoin(root, name string) (string, error) {
 	}
 	return joined, nil
 }
+
+func ensureNoSymlinkPath(root, target string, allowMissingLeaf bool) error {
+	rel, err := filepath.Rel(root, target)
+	if err != nil {
+		return fmt.Errorf("resolve relative path for %q: %w", target, err)
+	}
+	if rel == "." {
+		return nil
+	}
+	if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
+		return fmt.Errorf("refusing path outside extraction root %q", target)
+	}
+
+	parts := strings.Split(rel, string(filepath.Separator))
+	current := root
+	for i, part := range parts {
+		if part == "" || part == "." {
+			continue
+		}
+		current = filepath.Join(current, part)
+		info, statErr := os.Lstat(current)
+		if statErr != nil {
+			if errors.Is(statErr, os.ErrNotExist) {
+				if i == len(parts)-1 && allowMissingLeaf {
+					return nil
+				}
+				continue
+			}
+			return fmt.Errorf("inspect path %q: %w", current, statErr)
+		}
+		if info.Mode()&os.ModeSymlink != 0 {
+			return fmt.Errorf("refusing archive entry that traverses symlink path component %q", current)
+		}
+	}
+	return nil
+}
+
+func validateSymlinkTarget(root, targetPath, linkName string) error {
+	if strings.TrimSpace(linkName) == "" {
+		return fmt.Errorf("refusing symlink %q with empty target", targetPath)
+	}
+	resolved := filepath.Clean(filepath.Join(filepath.Dir(targetPath), linkName))
+	rel, err := filepath.Rel(root, resolved)
+	if err != nil {
+		return fmt.Errorf("resolve symlink target for %q: %w", targetPath, err)
+	}
+	if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
+		return fmt.Errorf("refusing symlink %q -> %q that resolves outside extraction root", targetPath, linkName)
+	}
+	return nil
+}