llmgw/internal/bundle/bundle_test.go

package bundle

import (
	"archive/tar"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
)

func hexSum(b []byte) string {
	s := sha256.Sum256(b)
	return hex.EncodeToString(s[:])
}

// internet is everything a bundle is made from: a release server, GitHub
// and Hugging Face, all fakes.
type internet struct {
	srv      *httptest.Server
	files    map[string][]byte
	digests  map[string]string // GitHub's, by asset
	priv     string
	weights  []byte
	hfDigest string
}

func newInternet(t *testing.T) *internet {
	t.Helper()
	pub, priv, err := release.GenerateKey()
	if err != nil {
		t.Fatal(err)
	}
	t.Setenv("LLMGW_RELEASE_KEY", pub)
	in := &internet{files: map[string][]byte{}, digests: map[string]string{}, priv: priv}

	// llmgw v1.2.3 for both architectures, signed.
	var sums strings.Builder
	for _, arch := range []string{"amd64", "arm64"} {
		name := release.AssetName("v1.2.3", "linux", arch)
		body := []byte("llmgw for " + arch)
		in.files["/releases/v1.2.3/"+name] = body
		sums.WriteString(hexSum(body) + "  " + name + "\n")
	}
	in.files["/releases/v1.2.3/checksums.txt"] = []byte(sums.String())
	sig, _ := release.Sign(priv, []byte(sums.String()))
	in.files["/releases/v1.2.3/"+release.SignatureName] = []byte(sig)

	// llama.cpp b1: every build there is for amd64.
	builds, err := provision.Builds("llamacpp", "b1", "amd64")
	if err != nil {
		t.Fatal(err)
	}
	for _, b := range builds {
		for _, a := range b.Assets {
			body := []byte("archive " + a)
			in.files["/gh/ggml-org/llama.cpp/releases/download/b1/"+a] = body
			in.digests[a] = "sha256:" + hexSum(body)
		}
	}

	// A model on Hugging Face, with the checksum the registry publishes.
	in.weights = []byte("GGUF weights")
	in.hfDigest = hexSum(in.weights)
	in.files["/Org/Model-GGUF/resolve/abc123/model-q4.gguf"] = in.weights

	in.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasPrefix(r.URL.Path, "/api/repos/ggml-org/llama.cpp/releases/tags/b1"):
			var out struct {
				Assets []map[string]string `json:"assets"`
			}
			for name, d := range in.digests {
				out.Assets = append(out.Assets, map[string]string{"name": name, "digest": d})
			}
			_ = json.NewEncoder(w).Encode(out)
		case r.URL.Path == "/api/models/Org/Model-GGUF":
			_, _ = w.Write([]byte(`{"id":"Org/Model-GGUF","sha":"abc123"}`))
		case r.URL.Path == "/api/models/Org/Model-GGUF/tree/abc123":
			_ = json.NewEncoder(w).Encode([]map[string]any{
				{"type": "file", "path": "model-q4.gguf", "size": 133,
					"lfs": map[string]any{"oid": in.hfDigest, "size": len(in.weights)}},
			})
		default:
			b, ok := in.files[r.URL.Path]
			if !ok {
				http.NotFound(w, r)
				return
			}
			_, _ = w.Write(b)
		}
	}))
	t.Cleanup(in.srv.Close)
	return in
}

func (in *internet) options(out string) CreateOptions {
	return CreateOptions{
		Out: out, Release: "v1.2.3", Arches: []string{"amd64"},
		Releases:  in.srv.URL + "/releases",
		Engines:   []EngineSpec{{Engine: "llamacpp", Version: "b1", Variants: []string{"cpu", "cuda-12.8"}}},
		Models:    []ModelSpec{{Name: "starter", Version: "2026-10", Source: "hf:Org/Model-GGUF/model-q4.gguf"}},
		HTTP:      in.srv.Client(),
		GitHubAPI: in.srv.URL + "/api", GitHubDownload: in.srv.URL + "/gh", Registry: in.srv.URL,
	}
}

func TestABundleCarriesWhatTheClusterWouldFetchAndAMirrorTakesItChecked(t *testing.T) {
	in := newInternet(t)
	out := filepath.Join(t.TempDir(), "bundle")
	m, digest, err := Create(t.Context(), in.options(out))
	if err != nil {
		t.Fatal(err)
	}
	var paths []string
	for _, f := range m.Files {
		paths = append(paths, f.Path)
	}
	for _, want := range []string{
		"releases/v1.2.3/checksums.txt", "releases/v1.2.3/" + release.SignatureName,
		"releases/v1.2.3/llmgw-v1.2.3-linux-amd64",
		"engines/llamacpp/b1/llama-b1-bin-ubuntu-x64.tar.gz",
		"engines/llamacpp/b1/llama-b1-bin-ubuntu-cuda-12.8-x64.tar.gz",
		"engines/llamacpp/b1/cudart-llama-b1-bin-ubuntu-cuda-12.8-x64.tar.gz",
		"models/starter/2026-10/model-q4.gguf",
	} {
		if !strings.Contains(strings.Join(paths, "\n"), want) {
			t.Errorf("the bundle lacks %s:\n%s", want, strings.Join(paths, "\n"))
		}
	}
	if strings.Contains(strings.Join(paths, "\n"), "vulkan") {
		t.Error("a build not asked for was carried")
	}

	// Carried as one file, and imported from it as from the directory.
	tarball := out + ".tar"
	if err := Pack(out, tarball); err != nil {
		t.Fatal(err)
	}
	for _, from := range []string{tarball, out} {
		root := t.TempDir()
		res, err := Import(from, ImportOptions{Mirror: root, SHA256: digest})
		if err != nil {
			t.Fatalf("importing %s: %v", filepath.Base(from), err)
		}
		if res.Written != len(m.Files) || res.Digest != digest {
			t.Errorf("wrote %d of %d files", res.Written, len(m.Files))
		}
		bin, _ := os.ReadFile(filepath.Join(root, "_llmgw/releases/v1.2.3/llmgw-v1.2.3-linux-amd64"))
		if string(bin) != "llmgw for amd64" {
			t.Errorf("the release is not where releases are fetched from: %q", bin)
		}
		if _, err := os.Stat(filepath.Join(root, "_llmgw/engines/llamacpp/b1/llama-b1-bin-ubuntu-x64.tar.gz")); err != nil {
			t.Error(err)
		}
		// The model is the mirror's like any other, with its checksum.
		man, err := mirror.ManifestFor(root, "starter", "2026-10")
		if err != nil || len(man.Files) != 1 || man.Files[0].SHA256 != in.hfDigest {
			t.Errorf("the model is not offered with its checksum: %+v %v", man, err)
		}
		// And the mirror lists nothing of llmgw's own as a model.
		entries, _ := mirror.List(root)
		if len(entries) != 1 {
			t.Errorf("the mirror lists %v", entries)
		}

		// Again: everything is there and the same, so nothing is written.
		res, err = Import(from, ImportOptions{Mirror: root, SHA256: digest})
		if err != nil || res.Written != 0 || res.Kept != len(m.Files) {
			t.Errorf("a second import: %+v %v", res, err)
		}
	}

	// Made again into the same directory, a file already there and right is
	// not downloaded: only the checksums are asked for again.
	for p := range in.files {
		if !strings.HasSuffix(p, "checksums.txt") && !strings.HasSuffix(p, release.SignatureName) {
			delete(in.files, p)
		}
	}
	_, again, err := Create(t.Context(), in.options(out))
	if err != nil {
		t.Fatalf("made again from what was there: %v", err)
	}
	var m2 Manifest
	raw, _ := os.ReadFile(filepath.Join(out, ManifestName))
	_ = json.Unmarshal(raw, &m2)
	if again != Digest(raw) || len(m2.Files) != len(m.Files) {
		t.Errorf("made again: %d files, digest %s", len(m2.Files), again)
	}
}

func TestTheControlPlaneTrustsTheBundlesEngineChecksumsAndNodesAreToldThem(t *testing.T) {
	in := newInternet(t)
	out := filepath.Join(t.TempDir(), "bundle")
	_, digest, err := Create(t.Context(), in.options(out))
	if err != nil {
		t.Fatal(err)
	}
	file := filepath.Join(t.TempDir(), "offline.sha256")
	if _, _, err := Trust(out, file, "", false); err == nil {
		t.Fatal("a bundle was trusted with no digest to check it by")
	}
	if _, _, err := Trust(out, file, strings.Repeat("0", 64), false); err == nil {
		t.Fatal("a bundle was trusted against the wrong digest")
	}
	if _, n, err := Trust(out, file, digest, false); err != nil || n != 3 {
		t.Fatalf("trusted %d: %v", n, err)
	}

	tr := NewTrusted(file)
	got := tr.For("llamacpp", "b1")
	if got["llama-b1-bin-ubuntu-x64.tar.gz"] != strings.TrimPrefix(in.digests["llama-b1-bin-ubuntu-x64.tar.gz"], "sha256:") {
		t.Errorf("the control plane holds %v", got)
	}
	if len(tr.For("llamacpp", "b2")) != 0 {
		t.Error("checksums were given for a version no bundle carried")
	}

	// A second bundle adds to what is trusted, and the control plane sees
	// it without a restart.
	m2 := Manifest{Format: Format, Files: []File{{Path: "engines/llamacpp/b2/llama-b2-bin-ubuntu-x64.tar.gz", Bytes: 1, SHA256: strings.Repeat("b", 64)}}}
	if n, err := TrustManifest(m2, file); err != nil || n != 4 {
		t.Fatalf("trusted %d: %v", n, err)
	}
	if tr.For("llamacpp", "b2")["llama-b2-bin-ubuntu-x64.tar.gz"] != strings.Repeat("b", 64) {
		t.Error("a bundle trusted while running was not seen")
	}
	if len(tr.For("llamacpp", "b1")) != 3 {
		t.Error("the first bundle's checksums were lost")
	}
}

func TestADamagedBundlePutsNothingInTheMirror(t *testing.T) {
	in := newInternet(t)
	out := filepath.Join(t.TempDir(), "bundle")
	_, digest, err := Create(t.Context(), in.options(out))
	if err != nil {
		t.Fatal(err)
	}
	// The wrong digest, and none at all, are refused before anything is read.
	root := t.TempDir()
	if _, err := Import(out, ImportOptions{Mirror: root, SHA256: strings.Repeat("0", 64)}); err == nil {
		t.Error("a bundle not matching its digest was imported")
	}
	if _, err := Import(out, ImportOptions{Mirror: root}); err == nil {
		t.Error("a bundle was imported with no digest")
	}
	if _, err := Import(out, ImportOptions{Mirror: root, Unchecked: true}); err != nil {
		t.Errorf("a bundle taken unchecked, as asked: %v", err)
	}

	// An engine changed on the way.
	engine := filepath.Join(out, "engines/llamacpp/b1/llama-b1-bin-ubuntu-x64.tar.gz")
	good, _ := os.ReadFile(engine)
	_ = os.WriteFile(engine, []byte("swapped"), 0o644)
	root = t.TempDir()
	if _, err := Import(out, ImportOptions{Mirror: root, SHA256: digest}); err == nil || !strings.Contains(err.Error(), "damaged") {
		t.Errorf("a changed engine: %v", err)
	}
	if entries := walk(t, root); len(entries) > 0 {
		t.Errorf("a damaged bundle left %v in the mirror", entries)
	}
	_ = os.WriteFile(engine, good, 0o644)

	// A release whose checksums the project did not sign, with a manifest
	// made to agree: what the manifest says is not enough for a binary.
	_, other, _ := release.GenerateKey()
	sums, _ := os.ReadFile(filepath.Join(out, "releases/v1.2.3/checksums.txt"))
	forged, _ := release.Sign(other, sums)
	_ = os.WriteFile(filepath.Join(out, "releases/v1.2.3/"+release.SignatureName), []byte(forged), 0o644)
	raw := remanifest(t, out)
	root = t.TempDir()
	if _, err := Import(out, ImportOptions{Mirror: root, SHA256: Digest(raw)}); err == nil || !strings.Contains(err.Error(), "not signed") {
		t.Errorf("a release not signed by the project: %v", err)
	}
	if entries := walk(t, root); len(entries) > 0 {
		t.Errorf("an unsigned release left %v in the mirror", entries)
	}
}

func TestAVersionAlreadyInTheMirrorIsNeverReplaced(t *testing.T) {
	in := newInternet(t)
	out := filepath.Join(t.TempDir(), "bundle")
	_, digest, err := Create(t.Context(), in.options(out))
	if err != nil {
		t.Fatal(err)
	}
	root := t.TempDir()
	_ = os.MkdirAll(filepath.Join(root, "starter", "2026-10"), 0o755)
	_ = os.WriteFile(filepath.Join(root, "starter", "2026-10", "model-q4.gguf"), []byte("other weights"), 0o644)
	if _, err := Import(out, ImportOptions{Mirror: root, SHA256: digest}); err == nil || !strings.Contains(err.Error(), "never replaced") {
		t.Fatalf("a different file under the same version: %v", err)
	}
	if b, _ := os.ReadFile(filepath.Join(root, "starter", "2026-10", "model-q4.gguf")); string(b) != "other weights" {
		t.Error("the mirror's file was replaced")
	}
}

func TestABundleCannotWriteOutsideTheMirror(t *testing.T) {
	for _, p := range []string{
		"../etc/cron.d/x", "models/../../x/y/z", "models/a/b/../c", "/abs/olute/path/x",
		"engines/llamacpp/b1", "releases/latest/llmgw", "models/_llmgw/releases/x",
		"models/.hidden/v/f", "other/a/b/c",
	} {
		if _, err := Destination(p); err == nil {
			t.Errorf("%q was given a destination", p)
		}
	}
	if d, err := Destination("models/m/v/f.gguf"); err != nil || d != "m/v/f.gguf" {
		t.Errorf("a model: %q %v", d, err)
	}

	// A tar carrying a file its manifest does not list is refused.
	dir := t.TempDir()
	raw, _ := json.Marshal(Manifest{Format: Format, Files: []File{}})
	tarball := filepath.Join(dir, "b.tar")
	f, _ := os.Create(tarball)
	tw := tar.NewWriter(f)
	_ = tw.WriteHeader(&tar.Header{Name: ManifestName, Mode: 0o644, Size: int64(len(raw)), Typeflag: tar.TypeReg})
	_, _ = tw.Write(raw)
	_ = tw.WriteHeader(&tar.Header{Name: "models/m/v/extra", Mode: 0o644, Size: 1, Typeflag: tar.TypeReg})
	_, _ = tw.Write([]byte("x"))
	_ = tw.Close()
	_ = f.Close()
	if _, err := Import(tarball, ImportOptions{Mirror: t.TempDir(), SHA256: Digest(raw)}); err == nil {
		t.Error("a file the manifest does not list was imported")
	}
}

// remanifest rewrites a bundle's manifest to agree with its files, as
// someone changing a bundle on the way would.
func remanifest(t *testing.T, dir string) []byte {
	t.Helper()
	raw, _ := os.ReadFile(filepath.Join(dir, ManifestName))
	var m Manifest
	_ = json.Unmarshal(raw, &m)
	for i, f := range m.Files {
		b, _ := os.ReadFile(filepath.Join(dir, f.Path))
		m.Files[i].SHA256, m.Files[i].Bytes = hexSum(b), int64(len(b))
	}
	raw, _ = json.MarshalIndent(m, "", "  ")
	_ = os.WriteFile(filepath.Join(dir, ManifestName), raw, 0o644)
	return raw
}

func walk(t *testing.T, root string) []string {
	t.Helper()
	var out []string
	_ = filepath.Walk(root, func(p string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() {
			out = append(out, p)
		}
		return nil
	})
	return out
}