banger/internal/daemon/images_pull.go

package daemon

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"banger/internal/api"
	"banger/internal/daemon/imagemgr"
	"banger/internal/imagepull"
	"banger/internal/model"

	"github.com/google/go-containerregistry/pkg/name"
)

// minPullExt4Size keeps the floor consistent with imagepull.MinExt4Size
// when the caller doesn't override --size and the OCI tree is tiny.
const minPullExt4Size int64 = 1 << 30 // 1 GiB

// PullImage downloads an OCI image, flattens it into an ext4 rootfs, and
// registers it as a managed banger image. Kernel info comes via --kernel-ref
// or direct paths, mirroring RegisterImage.
//
// The pulled rootfs's file ownership is the runner's uid/gid (Phase A v1
// limitation; see internal/imagepull). The image is suitable as input to
// `image build --from-image` but is not directly bootable until a future
// fixup pass lands.
func (d *Daemon) PullImage(ctx context.Context, params api.ImagePullParams) (image model.Image, err error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()

	ref := strings.TrimSpace(params.Ref)
	if ref == "" {
		return model.Image{}, errors.New("oci reference is required")
	}
	parsed, err := name.ParseReference(ref)
	if err != nil {
		return model.Image{}, fmt.Errorf("parse oci ref %q: %w", ref, err)
	}

	imgName := strings.TrimSpace(params.Name)
	if imgName == "" {
		imgName = defaultImageNameFromRef(parsed)
		if imgName == "" {
			return model.Image{}, errors.New("could not derive image name from ref; pass --name")
		}
	}
	if existing, lookupErr := d.store.GetImageByName(ctx, imgName); lookupErr == nil {
		return model.Image{}, fmt.Errorf("image %q already exists (id=%s); pick a different --name or delete it first", imgName, existing.ID)
	}

	kernelPath, initrdPath, modulesDir, err := d.resolveKernelInputs(params.KernelRef, params.KernelPath, params.InitrdPath, params.ModulesDir)
	if err != nil {
		return model.Image{}, err
	}
	if err := imagemgr.ValidateKernelPaths(kernelPath, initrdPath, modulesDir); err != nil {
		return model.Image{}, err
	}

	id, err := model.NewID()
	if err != nil {
		return model.Image{}, err
	}
	finalDir := filepath.Join(d.layout.ImagesDir, id)
	stagingDir := finalDir + ".staging"
	if err := os.MkdirAll(stagingDir, 0o755); err != nil {
		return model.Image{}, err
	}
	cleanupStaging := true
	defer func() {
		if cleanupStaging {
			_ = os.RemoveAll(stagingDir)
		}
	}()

	// Extract OCI layers into a working tree under TempDir so the
	// state filesystem doesn't temporarily double in size.
	rootfsTree, err := os.MkdirTemp("", "banger-pull-")
	if err != nil {
		return model.Image{}, err
	}
	defer os.RemoveAll(rootfsTree)

	meta, err := d.runPullAndFlatten(ctx, ref, d.layout.OCICacheDir, rootfsTree)
	if err != nil {
		return model.Image{}, fmt.Errorf("pull oci image: %w", err)
	}

	sizeBytes := params.SizeBytes
	if sizeBytes <= 0 {
		treeSize, err := dirSizeBytes(rootfsTree)
		if err != nil {
			return model.Image{}, fmt.Errorf("size oci tree: %w", err)
		}
		sizeBytes = treeSize + treeSize/4 // +25% headroom
		if sizeBytes < minPullExt4Size {
			sizeBytes = minPullExt4Size
		}
	}

	rootfsExt4 := filepath.Join(stagingDir, "rootfs.ext4")
	if err := imagepull.BuildExt4(ctx, d.runner, rootfsTree, rootfsExt4, sizeBytes); err != nil {
		return model.Image{}, fmt.Errorf("build rootfs ext4: %w", err)
	}
	if err := imagepull.ApplyOwnership(ctx, d.runner, rootfsExt4, meta); err != nil {
		return model.Image{}, fmt.Errorf("apply ownership: %w", err)
	}

	stagedKernel, stagedInitrd, stagedModules, err := imagemgr.StageBootArtifacts(ctx, d.runner, stagingDir, kernelPath, initrdPath, modulesDir)
	if err != nil {
		return model.Image{}, fmt.Errorf("stage boot artifacts: %w", err)
	}

	if err := os.Rename(stagingDir, finalDir); err != nil {
		return model.Image{}, fmt.Errorf("publish artifact dir: %w", err)
	}
	cleanupStaging = false

	now := model.Now()
	image = model.Image{
		ID:          id,
		Name:        imgName,
		Managed:     true,
		ArtifactDir: finalDir,
		RootfsPath:  filepath.Join(finalDir, filepath.Base(rootfsExt4)),
		KernelPath:  rebaseUnder(stagedKernel, stagingDir, finalDir),
		InitrdPath:  rebaseUnder(stagedInitrd, stagingDir, finalDir),
		ModulesDir:  rebaseUnder(stagedModules, stagingDir, finalDir),
		CreatedAt:   now,
		UpdatedAt:   now,
	}
	if err := d.store.UpsertImage(ctx, image); err != nil {
		_ = os.RemoveAll(finalDir)
		return model.Image{}, err
	}
	return image, nil
}

// runPullAndFlatten is the seam tests substitute. nil → real implementation.
func (d *Daemon) runPullAndFlatten(ctx context.Context, ref, cacheDir, destDir string) (imagepull.Metadata, error) {
	if d.pullAndFlatten != nil {
		return d.pullAndFlatten(ctx, ref, cacheDir, destDir)
	}
	pulled, err := imagepull.Pull(ctx, ref, cacheDir)
	if err != nil {
		return imagepull.Metadata{}, err
	}
	return imagepull.Flatten(ctx, pulled, destDir)
}

// nameSanitize keeps lowercase alphanumerics + hyphens, collapses runs.
var nameSanitizeRE = regexp.MustCompile(`[^a-z0-9]+`)

// defaultImageNameFromRef derives a friendly name like "debian-bookworm"
// from "docker.io/library/debian:bookworm". Returns "" if it can't.
func defaultImageNameFromRef(ref name.Reference) string {
	repo := ref.Context().RepositoryStr() // e.g. library/debian
	parts := strings.Split(repo, "/")
	base := parts[len(parts)-1]

	suffix := ""
	switch r := ref.(type) {
	case name.Tag:
		if t := r.TagStr(); t != "" && t != "latest" {
			suffix = "-" + t
		}
	case name.Digest:
		// take the first 12 hex chars after sha256:
		d := r.DigestStr()
		if i := strings.Index(d, ":"); i >= 0 && len(d) >= i+13 {
			suffix = "-" + d[i+1:i+13]
		}
	}

	out := nameSanitizeRE.ReplaceAllString(strings.ToLower(base+suffix), "-")
	out = strings.Trim(out, "-")
	return out
}

// rebaseUnder rewrites a path that points inside oldRoot to point inside
// newRoot. Empty input returns empty (kept by StageBootArtifacts when an
// optional artifact is absent).
func rebaseUnder(path, oldRoot, newRoot string) string {
	if path == "" {
		return ""
	}
	if rel, err := filepath.Rel(oldRoot, path); err == nil && !strings.HasPrefix(rel, "..") {
		return filepath.Join(newRoot, rel)
	}
	return path
}

// dirSizeBytes returns the sum of regular-file sizes under root, following
// no symlinks (lstat). Suitable for sizing an ext4 image.
func dirSizeBytes(root string) (int64, error) {
	var total int64
	err := filepath.WalkDir(root, func(_ string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.Type().IsRegular() {
			return nil
		}
		info, err := d.Info()
		if err != nil {
			return err
		}
		total += info.Size()
		return nil
	})
	return total, err
}