banger/internal/runtimebundle/bundle.go

package runtimebundle

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"strings"

	toml "github.com/pelletier/go-toml"
)

type Manifest struct {
	Version       string         `toml:"version"`
	URL           string         `toml:"url"`
	SHA256        string         `toml:"sha256"`
	BundleRoot    string         `toml:"bundle_root"`
	RequiredPaths []string       `toml:"required_paths"`
	BundleMeta    BundleMetadata `toml:"bundle_metadata"`
}

type BundleMetadata struct {
	FirecrackerBin      string `json:"firecracker_bin" toml:"firecracker_bin"`
	SSHKeyPath          string `json:"ssh_key_path" toml:"ssh_key_path"`
	NamegenPath         string `json:"namegen_path" toml:"namegen_path"`
	CustomizeScript     string `json:"customize_script" toml:"customize_script"`
	VSockPingHelperPath string `json:"vsock_ping_helper_path" toml:"vsock_ping_helper_path"`
	DefaultPackages     string `json:"default_packages_file" toml:"default_packages_file"`
	DefaultRootfs       string `json:"default_rootfs" toml:"default_rootfs"`
	DefaultBaseRootfs   string `json:"default_base_rootfs,omitempty" toml:"default_base_rootfs"`
	DefaultWorkSeed     string `json:"default_work_seed,omitempty" toml:"default_work_seed"`
	DefaultKernel       string `json:"default_kernel" toml:"default_kernel"`
	DefaultInitrd       string `json:"default_initrd,omitempty" toml:"default_initrd"`
	DefaultModulesDir   string `json:"default_modules_dir,omitempty" toml:"default_modules_dir"`
}

const BundleMetadataFile = "bundle.json"

func LoadManifest(path string) (Manifest, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return Manifest{}, err
	}
	var manifest Manifest
	if err := toml.Unmarshal(data, &manifest); err != nil {
		return Manifest{}, err
	}
	manifest.BundleRoot = strings.TrimSpace(manifest.BundleRoot)
	manifest.URL = strings.TrimSpace(manifest.URL)
	manifest.SHA256 = strings.ToLower(strings.TrimSpace(manifest.SHA256))
	manifest.BundleMeta = normalizeBundleMetadata(manifest.BundleMeta)
	for i, required := range manifest.RequiredPaths {
		manifest.RequiredPaths[i] = filepath.Clean(strings.TrimSpace(required))
	}
	sort.Strings(manifest.RequiredPaths)
	if len(manifest.RequiredPaths) == 0 {
		return Manifest{}, fmt.Errorf("runtime bundle manifest %s has no required_paths", path)
	}
	return manifest, nil
}

func Bootstrap(ctx context.Context, manifest Manifest, manifestPath, outDir string) error {
	if manifest.URL == "" {
		return fmt.Errorf("runtime bundle manifest %s has no url; point a local manifest copy at a staged or published runtime bundle archive", manifestPath)
	}
	if manifest.SHA256 == "" {
		return fmt.Errorf("runtime bundle manifest %s has no sha256; add the checksum for the staged or published runtime bundle archive", manifestPath)
	}
	manifestDir := filepath.Dir(manifestPath)
	parentDir := filepath.Dir(outDir)
	if err := os.MkdirAll(parentDir, 0o755); err != nil {
		return err
	}

	workDir, err := os.MkdirTemp(parentDir, ".runtime-bundle-*")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)

	archivePath := filepath.Join(workDir, "bundle.tar.gz")
	if err := downloadArchive(ctx, resolveSource(manifestDir, manifest.URL), archivePath); err != nil {
		return err
	}
	sum, err := fileSHA256(archivePath)
	if err != nil {
		return err
	}
	if sum != manifest.SHA256 {
		return fmt.Errorf("runtime bundle checksum mismatch: got %s want %s", sum, manifest.SHA256)
	}

	extractDir := filepath.Join(workDir, "extract")
	if err := extractTarGz(archivePath, extractDir); err != nil {
		return err
	}

	bundleDir := extractDir
	if manifest.BundleRoot != "" {
		bundleDir = filepath.Join(extractDir, manifest.BundleRoot)
	}
	if err := ValidateBundle(bundleDir, manifest.RequiredPaths); err != nil {
		return err
	}
	if _, err := LoadBundleMetadata(bundleDir); err != nil && !errors.Is(err, os.ErrNotExist) {
		return err
	}

	stageDir := filepath.Join(workDir, "stage")
	if err := os.Rename(bundleDir, stageDir); err != nil {
		return err
	}

	if err := os.RemoveAll(outDir); err != nil {
		return err
	}
	if err := os.Rename(stageDir, outDir); err != nil {
		return err
	}
	return nil
}

func ValidateBundle(bundleDir string, requiredPaths []string) error {
	for _, rel := range requiredPaths {
		if rel == "." || strings.HasPrefix(rel, "..") {
			return fmt.Errorf("invalid required bundle path: %s", rel)
		}
		if _, err := os.Stat(filepath.Join(bundleDir, rel)); err != nil {
			return fmt.Errorf("runtime bundle missing %s", rel)
		}
	}
	return nil
}

func Package(runtimeDir, outArchive string, manifest Manifest) (string, error) {
	if err := ValidateBundle(runtimeDir, manifest.RequiredPaths); err != nil {
		return "", err
	}
	metadata, err := metadataArchiveBytes(runtimeDir, manifest.BundleMeta)
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(filepath.Dir(outArchive), 0o755); err != nil {
		return "", err
	}
	file, err := os.Create(outArchive)
	if err != nil {
		return "", err
	}
	defer file.Close()

	hash := sha256.New()
	multi := io.MultiWriter(file, hash)
	gz := gzip.NewWriter(multi)
	defer gz.Close()
	tw := tar.NewWriter(gz)
	defer tw.Close()

	for _, rel := range manifest.RequiredPaths {
		if err := addPathToArchive(tw, runtimeDir, manifest.BundleRoot, rel); err != nil {
			return "", err
		}
	}
	if len(metadata) != 0 {
		if err := addBytesToArchive(tw, manifest.BundleRoot, BundleMetadataFile, metadata, 0o644); err != nil {
			return "", err
		}
	}
	if err := tw.Close(); err != nil {
		return "", err
	}
	if err := gz.Close(); err != nil {
		return "", err
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}

func LoadBundleMetadata(runtimeDir string) (BundleMetadata, error) {
	path := filepath.Join(runtimeDir, BundleMetadataFile)
	data, err := os.ReadFile(path)
	if err != nil {
		return BundleMetadata{}, err
	}
	var meta BundleMetadata
	if err := json.Unmarshal(data, &meta); err != nil {
		return BundleMetadata{}, fmt.Errorf("parse %s: %w", path, err)
	}
	meta = normalizeBundleMetadata(meta)
	if err := validateBundleMetadata(runtimeDir, meta); err != nil {
		return BundleMetadata{}, err
	}
	return meta, nil
}

func validateBundleMetadata(runtimeDir string, meta BundleMetadata) error {
	required := []struct {
		value string
		label string
	}{
		{meta.FirecrackerBin, "firecracker_bin"},
		{meta.SSHKeyPath, "ssh_key_path"},
		{meta.NamegenPath, "namegen_path"},
		{meta.CustomizeScript, "customize_script"},
		{meta.VSockPingHelperPath, "vsock_ping_helper_path"},
		{meta.DefaultPackages, "default_packages_file"},
		{meta.DefaultRootfs, "default_rootfs"},
		{meta.DefaultKernel, "default_kernel"},
	}
	for _, field := range required {
		if strings.TrimSpace(field.value) == "" {
			return fmt.Errorf("runtime bundle metadata missing %s", field.label)
		}
	}
	for _, field := range []struct {
		value    string
		label    string
		required bool
	}{
		{meta.FirecrackerBin, "firecracker_bin", true},
		{meta.SSHKeyPath, "ssh_key_path", true},
		{meta.NamegenPath, "namegen_path", true},
		{meta.CustomizeScript, "customize_script", true},
		{meta.VSockPingHelperPath, "vsock_ping_helper_path", true},
		{meta.DefaultPackages, "default_packages_file", true},
		{meta.DefaultRootfs, "default_rootfs", true},
		{meta.DefaultBaseRootfs, "default_base_rootfs", false},
		{meta.DefaultWorkSeed, "default_work_seed", false},
		{meta.DefaultKernel, "default_kernel", true},
		{meta.DefaultInitrd, "default_initrd", false},
		{meta.DefaultModulesDir, "default_modules_dir", false},
	} {
		if strings.TrimSpace(field.value) == "" {
			continue
		}
		resolved, err := resolveMetadataPath(runtimeDir, field.value)
		if err != nil {
			return fmt.Errorf("runtime bundle metadata %s: %w", field.label, err)
		}
		if _, err := os.Stat(resolved); err != nil {
			if field.required || !errors.Is(err, os.ErrNotExist) {
				return fmt.Errorf("runtime bundle metadata %s points to missing path %s", field.label, resolved)
			}
		}
	}
	return nil
}

func resolveMetadataPath(runtimeDir, rel string) (string, error) {
	rel = filepath.Clean(strings.TrimSpace(rel))
	if rel == "." || rel == "" || filepath.IsAbs(rel) || strings.HasPrefix(rel, "..") {
		return "", fmt.Errorf("invalid relative path %q", rel)
	}
	return filepath.Join(runtimeDir, rel), nil
}

func metadataArchiveBytes(runtimeDir string, meta BundleMetadata) ([]byte, error) {
	meta = normalizeBundleMetadata(meta)
	if strings.TrimSpace(meta.FirecrackerBin) == "" &&
		strings.TrimSpace(meta.SSHKeyPath) == "" &&
		strings.TrimSpace(meta.NamegenPath) == "" &&
		strings.TrimSpace(meta.CustomizeScript) == "" &&
		strings.TrimSpace(meta.VSockPingHelperPath) == "" &&
		strings.TrimSpace(meta.DefaultPackages) == "" &&
		strings.TrimSpace(meta.DefaultRootfs) == "" &&
		strings.TrimSpace(meta.DefaultBaseRootfs) == "" &&
		strings.TrimSpace(meta.DefaultWorkSeed) == "" &&
		strings.TrimSpace(meta.DefaultKernel) == "" &&
		strings.TrimSpace(meta.DefaultInitrd) == "" &&
		strings.TrimSpace(meta.DefaultModulesDir) == "" {
		return nil, nil
	}
	if err := validateBundleMetadata(runtimeDir, meta); err != nil {
		return nil, err
	}
	return json.MarshalIndent(meta, "", "  ")
}

func normalizeBundleMetadata(meta BundleMetadata) BundleMetadata {
	meta.FirecrackerBin = strings.TrimSpace(meta.FirecrackerBin)
	meta.SSHKeyPath = strings.TrimSpace(meta.SSHKeyPath)
	meta.NamegenPath = strings.TrimSpace(meta.NamegenPath)
	meta.CustomizeScript = strings.TrimSpace(meta.CustomizeScript)
	meta.VSockPingHelperPath = strings.TrimSpace(meta.VSockPingHelperPath)
	meta.DefaultPackages = strings.TrimSpace(meta.DefaultPackages)
	meta.DefaultRootfs = strings.TrimSpace(meta.DefaultRootfs)
	meta.DefaultBaseRootfs = strings.TrimSpace(meta.DefaultBaseRootfs)
	meta.DefaultWorkSeed = strings.TrimSpace(meta.DefaultWorkSeed)
	meta.DefaultKernel = strings.TrimSpace(meta.DefaultKernel)
	meta.DefaultInitrd = strings.TrimSpace(meta.DefaultInitrd)
	meta.DefaultModulesDir = strings.TrimSpace(meta.DefaultModulesDir)
	return meta
}

func addPathToArchive(tw *tar.Writer, runtimeDir, bundleRoot, rel string) error {
	srcPath := filepath.Join(runtimeDir, rel)
	info, err := os.Lstat(srcPath)
	if err != nil {
		return err
	}
	archiveName := rel
	if bundleRoot != "" {
		archiveName = filepath.Join(bundleRoot, rel)
	}
	if info.IsDir() {
		header, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}
		header.Name = filepath.ToSlash(archiveName) + "/"
		if err := tw.WriteHeader(header); err != nil {
			return err
		}
		entries, err := os.ReadDir(srcPath)
		if err != nil {
			return err
		}
		for _, entry := range entries {
			childRel := filepath.Join(rel, entry.Name())
			if err := addPathToArchive(tw, runtimeDir, bundleRoot, childRel); err != nil {
				return err
			}
		}
		return nil
	}

	header, err := tar.FileInfoHeader(info, "")
	if err != nil {
		return err
	}
	header.Name = filepath.ToSlash(archiveName)
	if err := tw.WriteHeader(header); err != nil {
		return err
	}
	file, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = io.Copy(tw, file)
	return err
}

func addBytesToArchive(tw *tar.Writer, bundleRoot, rel string, data []byte, mode int64) error {
	name := rel
	if bundleRoot != "" {
		name = filepath.Join(bundleRoot, rel)
	}
	header := &tar.Header{
		Name: filepath.ToSlash(name),
		Mode: mode,
		Size: int64(len(data)),
	}
	if err := tw.WriteHeader(header); err != nil {
		return err
	}
	_, err := tw.Write(data)
	return err
}

func resolveSource(manifestDir, source string) string {
	parsed, err := url.Parse(source)
	if err == nil && parsed.Scheme != "" {
		return source
	}
	if filepath.IsAbs(source) {
		return source
	}
	return filepath.Join(manifestDir, source)
}

func downloadArchive(ctx context.Context, source, dst string) error {
	switch {
	case strings.HasPrefix(source, "http://"), strings.HasPrefix(source, "https://"):
		req, err := http.NewRequestWithContext(ctx, http.MethodGet, source, nil)
		if err != nil {
			return err
		}
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		defer resp.Body.Close()
		if resp.StatusCode != http.StatusOK {
			return fmt.Errorf("download runtime bundle: %s", resp.Status)
		}
		return writeFileFromReader(dst, resp.Body)
	case strings.HasPrefix(source, "file://"):
		parsed, err := url.Parse(source)
		if err != nil {
			return err
		}
		return copyFile(parsed.Path, dst)
	default:
		return copyFile(source, dst)
	}
}

func writeFileFromReader(dst string, reader io.Reader) error {
	file, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer file.Close()
	_, err = io.Copy(file, reader)
	return err
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	return writeFileFromReader(dst, in)
}

func extractTarGz(archivePath, outDir string) error {
	if err := os.MkdirAll(outDir, 0o755); err != nil {
		return err
	}
	file, err := os.Open(archivePath)
	if err != nil {
		return err
	}
	defer file.Close()
	gz, err := gzip.NewReader(file)
	if err != nil {
		return err
	}
	defer gz.Close()
	tr := tar.NewReader(gz)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		name := filepath.Clean(header.Name)
		if name == "." || strings.HasPrefix(name, "..") || filepath.IsAbs(name) {
			return fmt.Errorf("invalid archive entry: %s", header.Name)
		}
		target := filepath.Join(outDir, name)
		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, os.FileMode(header.Mode)); err != nil {
				return err
			}
		case tar.TypeReg, tar.TypeRegA:
			if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
				return err
			}
			file, err := os.OpenFile(target, os.O_CREATE|os.O_RDWR|os.O_TRUNC, os.FileMode(header.Mode))
			if err != nil {
				return err
			}
			if _, err := io.Copy(file, tr); err != nil {
				file.Close()
				return err
			}
			if err := file.Close(); err != nil {
				return err
			}
		default:
			return fmt.Errorf("unsupported archive entry type: %s", header.Name)
		}
	}
}

func fileSHA256(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer file.Close()
	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}