banger/internal/cli/make_bundle_test.go

package cli

import (
	"archive/tar"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"

	"banger/internal/imagecat"

	"github.com/klauspost/compress/zstd"
)

func TestInternalMakeBundleFlagsExist(t *testing.T) {
	root := NewBangerCommand()
	internal, _, err := root.Find([]string{"internal"})
	if err != nil {
		t.Fatalf("find internal: %v", err)
	}
	mk, _, err := internal.Find([]string{"make-bundle"})
	if err != nil {
		t.Fatalf("find make-bundle: %v", err)
	}
	for _, name := range []string{"rootfs-tar", "name", "distro", "arch", "kernel-ref", "description", "size", "out"} {
		if mk.Flags().Lookup(name) == nil {
			t.Errorf("missing flag %q", name)
		}
	}
}

func TestMakeBundleRequiresName(t *testing.T) {
	cmd := NewBangerCommand()
	cmd.SetArgs([]string{"internal", "make-bundle", "--rootfs-tar", "some.tar", "--out", "out.tar.zst"})
	cmd.SetOut(&bytes.Buffer{})
	cmd.SetErr(&bytes.Buffer{})
	err := cmd.Execute()
	if err == nil || !strings.Contains(err.Error(), "image name is required") {
		t.Fatalf("execute error = %v, want image-name-required", err)
	}
}

func TestMakeBundleRequiresRootfsTar(t *testing.T) {
	cmd := NewBangerCommand()
	cmd.SetArgs([]string{"internal", "make-bundle", "--name", "x", "--out", "out.tar.zst"})
	cmd.SetOut(&bytes.Buffer{})
	cmd.SetErr(&bytes.Buffer{})
	err := cmd.Execute()
	if err == nil || !strings.Contains(err.Error(), "--rootfs-tar is required") {
		t.Fatalf("execute error = %v, want --rootfs-tar required", err)
	}
}

func TestMakeBundleRequiresOut(t *testing.T) {
	cmd := NewBangerCommand()
	cmd.SetArgs([]string{"internal", "make-bundle", "--name", "x", "--rootfs-tar", "-"})
	cmd.SetOut(&bytes.Buffer{})
	cmd.SetErr(&bytes.Buffer{})
	err := cmd.Execute()
	if err == nil || !strings.Contains(err.Error(), "--out is required") {
		t.Fatalf("execute error = %v, want --out required", err)
	}
}

func TestWriteBundleTarZstRoundTrip(t *testing.T) {
	stage := t.TempDir()
	rootfsContent := []byte("fake-rootfs-bytes")
	rootfsPath := filepath.Join(stage, "rootfs.ext4")
	if err := os.WriteFile(rootfsPath, rootfsContent, 0o644); err != nil {
		t.Fatal(err)
	}
	manifest := imagecat.Manifest{Name: "debian-bookworm", Distro: "debian"}
	manifestJSON, _ := json.Marshal(manifest)
	manifestPath := filepath.Join(stage, "manifest.json")
	if err := os.WriteFile(manifestPath, manifestJSON, 0o644); err != nil {
		t.Fatal(err)
	}

	bundlePath := filepath.Join(stage, "bundle.tar.zst")
	if err := writeBundleTarZst(bundlePath, rootfsPath, manifestPath); err != nil {
		t.Fatalf("writeBundleTarZst: %v", err)
	}

	// Decode and verify.
	raw, err := os.Open(bundlePath)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { raw.Close() })
	zr, err := zstd.NewReader(raw)
	if err != nil {
		t.Fatal(err)
	}
	tr := tar.NewReader(zr)
	got := map[string][]byte{}
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		b, _ := io.ReadAll(tr)
		got[hdr.Name] = b
	}
	if !bytes.Equal(got[imagecat.RootfsFilename], rootfsContent) {
		t.Errorf("rootfs mismatch: got %q want %q", got[imagecat.RootfsFilename], rootfsContent)
	}
	if !bytes.Equal(got[imagecat.ManifestFilename], manifestJSON) {
		t.Errorf("manifest mismatch: got %q want %q", got[imagecat.ManifestFilename], manifestJSON)
	}
}

func TestSha256HexFile(t *testing.T) {
	dir := t.TempDir()
	content := []byte("hello world")
	p := filepath.Join(dir, "f")
	if err := os.WriteFile(p, content, 0o644); err != nil {
		t.Fatal(err)
	}
	got, err := sha256HexFile(p)
	if err != nil {
		t.Fatal(err)
	}
	expected := sha256.Sum256(content)
	if got != hex.EncodeToString(expected[:]) {
		t.Fatalf("sha256 = %q, want %q", got, hex.EncodeToString(expected[:]))
	}
}

func TestDirSize(t *testing.T) {
	dir := t.TempDir()
	_ = os.MkdirAll(filepath.Join(dir, "sub"), 0o755)
	_ = os.WriteFile(filepath.Join(dir, "a"), []byte("abc"), 0o644)          // 3
	_ = os.WriteFile(filepath.Join(dir, "sub", "b"), []byte("defgh"), 0o644) // 5
	// Symlink must not be counted.
	_ = os.Symlink(filepath.Join(dir, "a"), filepath.Join(dir, "link"))
	n, err := dirSize(dir)
	if err != nil {
		t.Fatal(err)
	}
	if n != 8 {
		t.Fatalf("dirSize = %d, want 8", n)
	}
}

// TestMakeBundleEndToEnd exercises the full pipeline against a tiny
// synthesized rootfs tar. Skips if any external tool (mkfs.ext4 /
// debugfs) or the companion banger-vsock-agent binary is unavailable.
func TestMakeBundleEndToEnd(t *testing.T) {
	if _, err := exec.LookPath("mkfs.ext4"); err != nil {
		t.Skip("mkfs.ext4 not installed")
	}
	if _, err := exec.LookPath("debugfs"); err != nil {
		t.Skip("debugfs not installed")
	}
	// Build companion binary if the build tree doesn't already have one.
	buildDir := findBuildBinDir(t)
	if buildDir == "" {
		t.Skip("build/bin not found; run `make build` to enable this test")
	}
	if _, err := os.Stat(filepath.Join(buildDir, "banger-vsock-agent")); err != nil {
		t.Skip("banger-vsock-agent not in build/bin; run `make build`")
	}
	// Ensure the banger binary also exists so CompanionBinaryPath
	// resolves (it looks alongside the banger binary).
	if _, err := os.Stat(filepath.Join(buildDir, "banger")); err != nil {
		t.Skip("banger not in build/bin; run `make build`")
	}

	// Build a minimal rootfs tar: just /etc/os-release and /tmp (a dir).
	dir := t.TempDir()
	tarPath := filepath.Join(dir, "rootfs.tar")
	if err := writeMinimalTar(tarPath); err != nil {
		t.Fatal(err)
	}
	outPath := filepath.Join(dir, "bundle.tar.zst")

	// Invoke via the cobra command to cover arg handling too.
	cmd := NewBangerCommand()
	cmd.SetArgs([]string{
		"internal", "make-bundle",
		"--rootfs-tar", tarPath,
		"--name", "test-bundle",
		"--distro", "debian",
		"--arch", "x86_64",
		"--kernel-ref", "generic-6.12",
		"--size", "64M",
		"--out", outPath,
	})
	var stderr bytes.Buffer
	cmd.SetOut(&bytes.Buffer{})
	cmd.SetErr(&stderr)
	// paths.CompanionBinaryPath looks alongside the banger binary, but
	// the test binary lives elsewhere. Use the env override instead.
	t.Setenv("BANGER_VSOCK_AGENT_BIN", filepath.Join(buildDir, "banger-vsock-agent"))
	cmd.SetContext(context.Background())
	if err := cmd.Execute(); err != nil {
		t.Fatalf("execute: %v\nstderr:\n%s", err, stderr.String())
	}

	if stat, err := os.Stat(outPath); err != nil {
		t.Fatalf("output not written: %v", err)
	} else if stat.Size() < 1024 {
		t.Fatalf("output suspiciously small: %d bytes", stat.Size())
	}

	// Verify we can fetch-reparse it (mirror of imagecat.Fetch logic,
	// but reading straight from disk instead of HTTP).
	extractDir := t.TempDir()
	verifyBundle(t, outPath, extractDir)
}

// findBuildBinDir returns the absolute path to the project's build/bin,
// or "" if it can't be located. Walks up from CWD to find go.mod.
func findBuildBinDir(t *testing.T) string {
	t.Helper()
	cwd, err := os.Getwd()
	if err != nil {
		return ""
	}
	for d := cwd; d != "/" && d != "."; d = filepath.Dir(d) {
		if _, err := os.Stat(filepath.Join(d, "go.mod")); err == nil {
			return filepath.Join(d, "build", "bin")
		}
	}
	return ""
}

func writeMinimalTar(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	tw := tar.NewWriter(f)
	defer tw.Close()

	// /etc dir
	if err := tw.WriteHeader(&tar.Header{
		Name: "etc/", Typeflag: tar.TypeDir, Mode: 0o755, Uid: 0, Gid: 0,
	}); err != nil {
		return err
	}
	// /etc/os-release
	body := []byte(`ID=debian` + "\n" + `PRETTY_NAME="banger test"` + "\n")
	if err := tw.WriteHeader(&tar.Header{
		Name: "etc/os-release", Typeflag: tar.TypeReg, Mode: 0o644,
		Size: int64(len(body)), Uid: 0, Gid: 0,
	}); err != nil {
		return err
	}
	if _, err := tw.Write(body); err != nil {
		return err
	}
	// /tmp dir
	return tw.WriteHeader(&tar.Header{
		Name: "tmp/", Typeflag: tar.TypeDir, Mode: 0o1777, Uid: 0, Gid: 0,
	})
}

func verifyBundle(t *testing.T, bundlePath, extractDir string) {
	t.Helper()
	f, err := os.Open(bundlePath)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	zr, err := zstd.NewReader(f)
	if err != nil {
		t.Fatal(err)
	}
	defer zr.Close()
	tr := tar.NewReader(zr)
	seen := map[string]bool{}
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		dst := filepath.Join(extractDir, hdr.Name)
		if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil {
			t.Fatal(err)
		}
		out, err := os.Create(dst)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := io.Copy(out, tr); err != nil {
			t.Fatal(err)
		}
		out.Close()
		seen[hdr.Name] = true
	}
	if !seen[imagecat.RootfsFilename] || !seen[imagecat.ManifestFilename] {
		t.Fatalf("bundle missing expected files: seen=%v", seen)
	}
	manifestData, err := os.ReadFile(filepath.Join(extractDir, imagecat.ManifestFilename))
	if err != nil {
		t.Fatal(err)
	}
	var m imagecat.Manifest
	if err := json.Unmarshal(manifestData, &m); err != nil {
		t.Fatal(err)
	}
	if m.Name != "test-bundle" || m.KernelRef != "generic-6.12" || m.Distro != "debian" {
		t.Fatalf("manifest = %+v", m)
	}
}