banger/internal/daemon/images.go

package daemon

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"banger/internal/api"
	"banger/internal/imagepreset"
	"banger/internal/model"
	"banger/internal/system"
)

func (d *Daemon) BuildImage(ctx context.Context, params api.ImageBuildParams) (image model.Image, err error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()
	op := d.beginOperation("image.build")
	buildLogPath := ""
	defer func() {
		if err != nil {
			err = annotateLogPath(err, buildLogPath)
			op.fail(err, imageLogAttrs(image)...)
			return
		}
		op.done(imageLogAttrs(image)...)
	}()

	name := params.Name
	imageBuildStage(ctx, "resolve_image", "resolving image build inputs")
	if name == "" {
		name = fmt.Sprintf("image-%d", model.Now().Unix())
	}
	if _, err := d.FindImage(ctx, name); err == nil {
		return model.Image{}, fmt.Errorf("image name already exists: %s", name)
	}
	fromImage := strings.TrimSpace(params.FromImage)
	if fromImage == "" {
		return model.Image{}, fmt.Errorf("from-image is required")
	}
	baseImage, err := d.FindImage(ctx, fromImage)
	if err != nil {
		return model.Image{}, err
	}
	id, err := model.NewID()
	if err != nil {
		return model.Image{}, err
	}
	now := model.Now()
	artifactDir := filepath.Join(d.layout.ImagesDir, id)
	buildLogDir := filepath.Join(d.layout.StateDir, "image-build")
	if err := os.MkdirAll(buildLogDir, 0o755); err != nil {
		return model.Image{}, err
	}
	buildLogPath = filepath.Join(buildLogDir, id+".log")
	imageBuildSetLogPath(ctx, buildLogPath)
	logFile, err := os.OpenFile(buildLogPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o644)
	if err != nil {
		return model.Image{}, err
	}
	defer logFile.Close()
	stageDir, err := os.MkdirTemp(d.layout.ImagesDir, id+".build-")
	if err != nil {
		return model.Image{}, err
	}
	cleanupStage := true
	defer func() {
		if cleanupStage {
			_ = os.RemoveAll(stageDir)
		}
	}()
	rootfsPath := filepath.Join(stageDir, "rootfs.ext4")
	workSeedPath := filepath.Join(stageDir, "work-seed.ext4")
	kernelSource := firstNonEmpty(params.KernelPath, baseImage.KernelPath)
	initrdSource := firstNonEmpty(params.InitrdPath, baseImage.InitrdPath)
	modulesSource := firstNonEmpty(params.ModulesDir, baseImage.ModulesDir)
	if err := d.validateImageBuildPrereqs(ctx, baseImage.RootfsPath, kernelSource, initrdSource, modulesSource, params.Size); err != nil {
		return model.Image{}, err
	}
	kernelPath, initrdPath, modulesDir, err := stageManagedBootArtifacts(ctx, d.runner, stageDir, kernelSource, initrdSource, modulesSource)
	if err != nil {
		return model.Image{}, err
	}
	packages := imagepreset.DebianBasePackages()
	metadataPackages := imageBuildMetadataPackages(params.Docker)
	spec := imageBuildSpec{
		ID:            id,
		Name:          name,
		SourceRootfs:  baseImage.RootfsPath,
		RootfsPath:    rootfsPath,
		BuildLog:      logFile,
		KernelPath:    kernelPath,
		InitrdPath:    initrdPath,
		ModulesDir:    modulesDir,
		Packages:      packages,
		InstallDocker: params.Docker,
		Size:          params.Size,
	}
	op.stage("launch_builder", "build_log_path", buildLogPath, "artifact_dir", artifactDir, "from_image", baseImage.Name)
	imageBuildStage(ctx, "launch_builder", "building rootfs from base image")
	if err := d.runImageBuild(ctx, spec); err != nil {
		_ = logFile.Sync()
		return model.Image{}, err
	}
	imageBuildStage(ctx, "prepare_work_seed", "building reusable work seed")
	if err := system.BuildWorkSeedImage(ctx, d.runner, rootfsPath, workSeedPath); err != nil {
		_ = logFile.Sync()
		return model.Image{}, err
	}
	imageBuildStage(ctx, "seed_ssh", "seeding runtime SSH access")
	seededSSHPublicKeyFingerprint, err := d.seedAuthorizedKeyOnExt4Image(ctx, workSeedPath)
	if err != nil {
		_ = logFile.Sync()
		return model.Image{}, err
	}
	imageBuildStage(ctx, "write_metadata", "writing image metadata")
	if err := writePackagesMetadata(rootfsPath, metadataPackages); err != nil {
		_ = logFile.Sync()
		return model.Image{}, err
	}
	op.stage("activate_artifacts", "artifact_dir", artifactDir)
	if err := os.Rename(stageDir, artifactDir); err != nil {
		return model.Image{}, err
	}
	cleanupStage = false
	image = model.Image{
		ID:                            id,
		Name:                          name,
		Managed:                       true,
		ArtifactDir:                   artifactDir,
		RootfsPath:                    filepath.Join(artifactDir, "rootfs.ext4"),
		WorkSeedPath:                  filepath.Join(artifactDir, "work-seed.ext4"),
		KernelPath:                    filepath.Join(artifactDir, "kernel"),
		InitrdPath:                    stageOptionalArtifactPath(artifactDir, initrdPath, "initrd.img"),
		ModulesDir:                    stageOptionalArtifactPath(artifactDir, modulesDir, "modules"),
		BuildSize:                     params.Size,
		SeededSSHPublicKeyFingerprint: seededSSHPublicKeyFingerprint,
		Docker:                        params.Docker,
		CreatedAt:                     now,
		UpdatedAt:                     now,
	}
	imageBuildBindImage(ctx, image)
	if err := d.store.UpsertImage(ctx, image); err != nil {
		return model.Image{}, err
	}
	op.stage("persisted", "build_log_path", buildLogPath)
	imageBuildStage(ctx, "persisted", "image metadata saved")
	if d.logger != nil {
		d.logger.Info("image build log preserved", append(imageLogAttrs(image), "build_log_path", buildLogPath)...)
	}
	_ = logFile.Sync()
	return image, nil
}

func (d *Daemon) RegisterImage(ctx context.Context, params api.ImageRegisterParams) (image model.Image, err error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()

	name := strings.TrimSpace(params.Name)
	if name == "" {
		return model.Image{}, fmt.Errorf("image name is required")
	}

	rootfsPath := strings.TrimSpace(params.RootfsPath)
	if rootfsPath == "" {
		return model.Image{}, fmt.Errorf("rootfs path is required")
	}
	workSeedPath := strings.TrimSpace(params.WorkSeedPath)
	if workSeedPath == "" {
		candidate := system.WorkSeedPath(rootfsPath)
		if candidate != "" {
			if _, statErr := os.Stat(candidate); statErr == nil {
				workSeedPath = candidate
			}
		}
	}
	kernelPath := strings.TrimSpace(params.KernelPath)
	if kernelPath == "" {
		return model.Image{}, fmt.Errorf("kernel path is required")
	}
	initrdPath := strings.TrimSpace(params.InitrdPath)
	modulesDir := strings.TrimSpace(params.ModulesDir)

	if err := validateImageRegisterPaths(rootfsPath, workSeedPath, kernelPath, initrdPath, modulesDir); err != nil {
		return model.Image{}, err
	}

	now := model.Now()
	existing, lookupErr := d.store.GetImageByName(ctx, name)
	switch {
	case lookupErr == nil:
		if existing.Managed {
			return model.Image{}, fmt.Errorf("managed image %s cannot be updated via register", name)
		}
		image = existing
		image.RootfsPath = rootfsPath
		image.WorkSeedPath = workSeedPath
		image.KernelPath = kernelPath
		image.InitrdPath = initrdPath
		image.ModulesDir = modulesDir
		image.Docker = params.Docker
		image.UpdatedAt = now
	case errors.Is(lookupErr, sql.ErrNoRows):
		id, idErr := model.NewID()
		if idErr != nil {
			return model.Image{}, idErr
		}
		image = model.Image{
			ID:           id,
			Name:         name,
			Managed:      false,
			RootfsPath:   rootfsPath,
			WorkSeedPath: workSeedPath,
			KernelPath:   kernelPath,
			InitrdPath:   initrdPath,
			ModulesDir:   modulesDir,
			Docker:       params.Docker,
			CreatedAt:    now,
			UpdatedAt:    now,
		}
	default:
		return model.Image{}, lookupErr
	}

	if err := d.store.UpsertImage(ctx, image); err != nil {
		return model.Image{}, err
	}
	return image, nil
}

func (d *Daemon) PromoteImage(ctx context.Context, idOrName string) (image model.Image, err error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()

	op := d.beginOperation("image.promote")
	defer func() {
		if err != nil {
			op.fail(err, imageLogAttrs(image)...)
			return
		}
		op.done(imageLogAttrs(image)...)
	}()

	image, err = d.FindImage(ctx, idOrName)
	if err != nil {
		return model.Image{}, err
	}
	if image.Managed {
		return model.Image{}, fmt.Errorf("image %s is already managed", image.Name)
	}
	if err := validateImagePromotePaths(image.RootfsPath, image.KernelPath, image.InitrdPath, image.ModulesDir); err != nil {
		return model.Image{}, err
	}
	if strings.TrimSpace(d.layout.ImagesDir) == "" {
		return model.Image{}, errors.New("images dir is not configured")
	}
	if err := os.MkdirAll(d.layout.ImagesDir, 0o755); err != nil {
		return model.Image{}, err
	}

	artifactDir := filepath.Join(d.layout.ImagesDir, image.ID)
	if _, statErr := os.Stat(artifactDir); statErr == nil {
		return model.Image{}, fmt.Errorf("artifact dir already exists: %s", artifactDir)
	} else if !os.IsNotExist(statErr) {
		return model.Image{}, statErr
	}

	stageDir, err := os.MkdirTemp(d.layout.ImagesDir, image.ID+".promote-")
	if err != nil {
		return model.Image{}, err
	}
	cleanupStage := true
	defer func() {
		if cleanupStage {
			_ = os.RemoveAll(stageDir)
		}
	}()

	rootfsPath := filepath.Join(stageDir, "rootfs.ext4")
	op.stage("copy_rootfs", "source_rootfs_path", image.RootfsPath, "target_rootfs_path", rootfsPath)
	if err := system.CopyFilePreferClone(image.RootfsPath, rootfsPath); err != nil {
		return model.Image{}, err
	}

	workSeedPath := ""
	if image.WorkSeedPath != "" {
		if _, statErr := os.Stat(image.WorkSeedPath); statErr != nil {
			if os.IsNotExist(statErr) {
				op.stage("skip_missing_work_seed", "source_work_seed_path", image.WorkSeedPath)
				image.WorkSeedPath = ""
			} else {
				return model.Image{}, statErr
			}
		}
	}
	if image.WorkSeedPath != "" {
		workSeedPath = filepath.Join(stageDir, "work-seed.ext4")
		op.stage("copy_work_seed", "source_work_seed_path", image.WorkSeedPath, "target_work_seed_path", workSeedPath)
		if err := system.CopyFilePreferClone(image.WorkSeedPath, workSeedPath); err != nil {
			return model.Image{}, err
		}
		image.SeededSSHPublicKeyFingerprint, err = d.seedAuthorizedKeyOnExt4Image(ctx, workSeedPath)
		if err != nil {
			return model.Image{}, err
		}
	} else {
		image.SeededSSHPublicKeyFingerprint = ""
	}
	_, initrdPath, modulesDir, err := stageManagedBootArtifacts(ctx, d.runner, stageDir, image.KernelPath, image.InitrdPath, image.ModulesDir)
	if err != nil {
		return model.Image{}, err
	}

	op.stage("activate_artifacts", "artifact_dir", artifactDir)
	if err := os.Rename(stageDir, artifactDir); err != nil {
		return model.Image{}, err
	}
	cleanupStage = false

	image.Managed = true
	image.ArtifactDir = artifactDir
	image.RootfsPath = filepath.Join(artifactDir, "rootfs.ext4")
	if workSeedPath != "" {
		image.WorkSeedPath = filepath.Join(artifactDir, "work-seed.ext4")
	}
	image.KernelPath = filepath.Join(artifactDir, "kernel")
	image.InitrdPath = stageOptionalArtifactPath(artifactDir, initrdPath, "initrd.img")
	image.ModulesDir = stageOptionalArtifactPath(artifactDir, modulesDir, "modules")
	image.UpdatedAt = model.Now()
	if err := d.store.UpsertImage(ctx, image); err != nil {
		_ = os.RemoveAll(artifactDir)
		return model.Image{}, err
	}
	return image, nil
}

func validateImageRegisterPaths(rootfsPath, workSeedPath, kernelPath, initrdPath, modulesDir string) error {
	checks := system.NewPreflight()
	checks.RequireFile(rootfsPath, "rootfs image", `pass --rootfs <path>`)
	checks.RequireFile(kernelPath, "kernel image", `pass --kernel <path>`)
	if workSeedPath != "" {
		checks.RequireFile(workSeedPath, "work-seed image", `pass --work-seed <path> or rebuild the image with a work seed`)
	}
	if initrdPath != "" {
		checks.RequireFile(initrdPath, "initrd image", `pass --initrd <path>`)
	}
	if modulesDir != "" {
		checks.RequireDir(modulesDir, "kernel modules dir", `pass --modules <dir>`)
	}
	return checks.Err("image register failed")
}

func validateImagePromotePaths(rootfsPath, kernelPath, initrdPath, modulesDir string) error {
	checks := system.NewPreflight()
	checks.RequireFile(rootfsPath, "rootfs image", `re-register the image with a valid rootfs`)
	checks.RequireFile(kernelPath, "kernel image", `re-register the image with a valid kernel`)
	if initrdPath != "" {
		checks.RequireFile(initrdPath, "initrd image", `re-register the image with a valid initrd`)
	}
	if modulesDir != "" {
		checks.RequireDir(modulesDir, "kernel modules dir", `re-register the image with a valid modules dir`)
	}
	return checks.Err("image promote failed")
}

func writePackagesMetadata(rootfsPath string, packages []string) error {
	if rootfsPath == "" || len(packages) == 0 {
		return nil
	}
	metadataPath := rootfsPath + ".packages.sha256"
	return os.WriteFile(metadataPath, []byte(packagesHash(packages)+"\n"), 0o644)
}

func (d *Daemon) DeleteImage(ctx context.Context, idOrName string) (model.Image, error) {
	d.imageOpsMu.Lock()
	defer d.imageOpsMu.Unlock()

	image, err := d.FindImage(ctx, idOrName)
	if err != nil {
		return model.Image{}, err
	}
	vms, err := d.store.FindVMsUsingImage(ctx, image.ID)
	if err != nil {
		return model.Image{}, err
	}
	if len(vms) > 0 {
		return model.Image{}, fmt.Errorf("image %s is still referenced by %d VM(s)", image.Name, len(vms))
	}
	if err := d.store.DeleteImage(ctx, image.ID); err != nil {
		return model.Image{}, err
	}
	if image.Managed && image.ArtifactDir != "" {
		if err := os.RemoveAll(image.ArtifactDir); err != nil {
			return model.Image{}, err
		}
	}
	return image, nil
}

func stageManagedBootArtifacts(ctx context.Context, runner system.CommandRunner, artifactDir, kernelSource, initrdSource, modulesSource string) (string, string, string, error) {
	kernelPath := filepath.Join(artifactDir, "kernel")
	if err := system.CopyFilePreferClone(kernelSource, kernelPath); err != nil {
		return "", "", "", err
	}
	initrdPath := ""
	if strings.TrimSpace(initrdSource) != "" {
		initrdPath = filepath.Join(artifactDir, "initrd.img")
		if err := system.CopyFilePreferClone(initrdSource, initrdPath); err != nil {
			return "", "", "", err
		}
	}
	modulesDir := ""
	if strings.TrimSpace(modulesSource) != "" {
		modulesDir = filepath.Join(artifactDir, "modules")
		if err := os.MkdirAll(modulesDir, 0o755); err != nil {
			return "", "", "", err
		}
		if err := system.CopyDirContents(ctx, runner, modulesSource, modulesDir, false); err != nil {
			return "", "", "", err
		}
	}
	return kernelPath, initrdPath, modulesDir, nil
}

func imageBuildMetadataPackages(docker bool) []string {
	packages := imagepreset.DebianBasePackages()
	if docker {
		packages = append(packages, "#feature:docker")
	}
	return packages
}

func stageOptionalArtifactPath(artifactDir, stagedPath, name string) string {
	if strings.TrimSpace(stagedPath) == "" {
		return ""
	}
	return filepath.Join(artifactDir, name)
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return value
		}
	}
	return ""
}