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/imagecat"
	"banger/internal/imagepull"
	"banger/internal/model"
	"banger/internal/paths"

	"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 image and registers it as a managed banger
// image. Two paths:
//
//   - Bundle path: `ref` matches an entry in the embedded imagecat
//     catalog. The `.tar.zst` bundle is fetched, `rootfs.ext4` is
//     already flattened + ownership-fixed + agent-injected at build
//     time, so this path is strictly faster than the OCI one.
//   - OCI path: otherwise treat `ref` as an OCI reference, pull its
//     layers, flatten, fix ownership, inject agents.
//
// Kernel info falls back through: `params.KernelRef` → catalog entry's
// `kernel_ref` (bundle path only) → `params.Kernel/Initrd/ModulesDir`.
func (d *Daemon) PullImage(ctx context.Context, params api.ImagePullParams) (model.Image, error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()

	ref := strings.TrimSpace(params.Ref)
	if ref == "" {
		return model.Image{}, errors.New("reference is required")
	}

	catalog, err := imagecat.LoadEmbedded()
	if err != nil {
		return model.Image{}, fmt.Errorf("load image catalog: %w", err)
	}
	if entry, lookupErr := catalog.Lookup(ref); lookupErr == nil {
		return d.pullFromBundle(ctx, params, entry)
	}
	return d.pullFromOCI(ctx, params)
}

// pullFromOCI is the original OCI-registry-pull path. See PullImage for
// the intent.
func (d *Daemon) pullFromOCI(ctx context.Context, params api.ImagePullParams) (image model.Image, err error) {
	ref := strings.TrimSpace(params.Ref)
	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(ctx, 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 := d.runFinalizePulledRootfs(ctx, rootfsExt4, meta); err != nil {
		return model.Image{}, 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
}

// pullFromBundle is the imagecat-backed path: download a ready-to-boot
// bundle (rootfs.ext4 already flattened + ownership-fixed + agent-
// injected at build time), verify its sha256, and register the result
// as a managed image. No flatten / mkfs / debugfs work on the daemon
// host.
func (d *Daemon) pullFromBundle(ctx context.Context, params api.ImagePullParams, entry imagecat.CatEntry) (image model.Image, err error) {
	imgName := strings.TrimSpace(params.Name)
	if imgName == "" {
		imgName = entry.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)
	}

	// Kernel resolution precedence: params > catalog entry's kernel_ref.
	kernelRef := strings.TrimSpace(params.KernelRef)
	if kernelRef == "" && strings.TrimSpace(params.KernelPath) == "" {
		kernelRef = strings.TrimSpace(entry.KernelRef)
	}
	kernelPath, initrdPath, modulesDir, err := d.resolveKernelInputs(ctx, 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)
		}
	}()

	if _, err := d.runBundleFetch(ctx, stagingDir, entry); err != nil {
		return model.Image{}, fmt.Errorf("fetch bundle: %w", err)
	}
	// manifest.json is metadata we only need at fetch time; strip it
	// so the final artifact dir contains only boot-relevant files.
	_ = os.Remove(filepath.Join(stagingDir, imagecat.ManifestFilename))
	rootfsExt4 := filepath.Join(stagingDir, imagecat.RootfsFilename)

	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
}

// runBundleFetch is the seam tests substitute. nil → real implementation.
func (d *Daemon) runBundleFetch(ctx context.Context, destDir string, entry imagecat.CatEntry) (imagecat.Manifest, error) {
	if d.bundleFetch != nil {
		return d.bundleFetch(ctx, destDir, entry)
	}
	return imagecat.Fetch(ctx, nil, destDir, entry)
}

// 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)
}

// runFinalizePulledRootfs applies ownership fixup and injects banger's
// guest agents. Tests substitute via d.finalizePulledRootfs; nil →
// real implementation using debugfs + the companion vsock-agent
// binary resolved via paths.CompanionBinaryPath.
func (d *Daemon) runFinalizePulledRootfs(ctx context.Context, ext4File string, meta imagepull.Metadata) error {
	if d.finalizePulledRootfs != nil {
		return d.finalizePulledRootfs(ctx, ext4File, meta)
	}
	if err := imagepull.ApplyOwnership(ctx, d.runner, ext4File, meta); err != nil {
		return fmt.Errorf("apply ownership: %w", err)
	}
	vsockBin, err := paths.CompanionBinaryPath("banger-vsock-agent")
	if err != nil {
		return fmt.Errorf("locate vsock agent binary: %w", err)
	}
	if err := imagepull.InjectGuestAgents(ctx, d.runner, ext4File, imagepull.GuestAgentAssets{
		VsockAgentBin: vsockBin,
	}); err != nil {
		return fmt.Errorf("inject guest agents: %w", err)
	}
	return nil
}

// 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
}