llmgw/internal/bundle/trust.go

package bundle

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"maps"
	"os"
	"path"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"time"
)

// Trust records a bundle's engine checksums in the file the control plane
// reads them from, beside any recorded before, and says how many it holds
// now. Only the manifest is read: the checksums are what each node holds
// the files it fetches to, so the manifest must be the one that was made —
// which is what the digest says.
func Trust(p, file, digest string, unchecked bool) (Manifest, int, error) {
	src, m, err := open(p)
	if err != nil {
		return Manifest{}, 0, err
	}
	defer src.close() //nolint:errcheck // read only
	if err := checkDigest(src.manifest(), digest, unchecked); err != nil {
		return m, 0, err
	}
	n, err := TrustManifest(m, file)
	return m, n, err
}

// TrustManifest records a manifest's engine checksums, for one already
// checked by an import.
func TrustManifest(m Manifest, file string) (int, error) {
	sums, err := readTrusted(file)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return 0, err
	}
	if sums == nil {
		sums = map[string]string{}
	}
	for _, f := range m.Files {
		if strings.HasPrefix(f.Path, kindEngine+"/") {
			sums[Prefix+"/"+f.Path] = f.SHA256
		}
	}
	var b strings.Builder
	b.WriteString("# Engine files the nodes may install, as the offline bundles they came in\n" +
		"# gave them. Kept by 'llmgw bundle import' and 'llmgw bundle trust'.\n")
	for _, k := range slices.Sorted(maps.Keys(sums)) {
		fmt.Fprintf(&b, "%s  %s\n", sums[k], k)
	}
	if err := os.MkdirAll(filepath.Dir(file), 0o755); err != nil {
		return 0, err
	}
	tmp := file + partSuffix
	if err := os.WriteFile(tmp, []byte(b.String()), 0o644); err != nil { //nolint:gosec // checksums, for anyone to read
		return 0, err
	}
	return len(sums), os.Rename(tmp, file)
}

// readTrusted reads a checksums file: "<sha256>  <path>" a line, as
// sha256sum writes them, with paths as the mirror has the files.
func readTrusted(file string) (map[string]string, error) {
	b, err := os.ReadFile(file) //nolint:gosec // the control plane's own file
	if err != nil {
		return nil, err
	}
	out := map[string]string{}
	sc := bufio.NewScanner(bytes.NewReader(b))
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		f := strings.Fields(line)
		if len(f) != 2 || len(f[0]) != 64 {
			return nil, fmt.Errorf("%s: %q is not a checksum and a file", file, line)
		}
		out[strings.TrimPrefix(f[1], "*")] = f[0]
	}
	return out, sc.Err()
}

// Trusted is the engine checksums a control plane holds, read from their
// file again whenever it changes — so a bundle imported while the control
// plane runs is offered to the next node that asks.
type Trusted struct {
	file string

	mu   sync.Mutex
	mod  time.Time
	size int64
	sums map[string]string
	err  error
}

// NewTrusted reads checksums from file as they are asked for.
func NewTrusted(file string) *Trusted { return &Trusted{file: file} }

// For is an engine version's files and their checksums, by file name; none
// when no bundle imported here carried it.
func (t *Trusted) For(engine, version string) map[string]string {
	t.mu.Lock()
	defer t.mu.Unlock()
	if info, err := os.Stat(t.file); err != nil {
		t.sums, t.err, t.mod, t.size = nil, err, time.Time{}, 0
	} else if !info.ModTime().Equal(t.mod) || info.Size() != t.size {
		t.sums, t.err = readTrusted(t.file)
		t.mod, t.size = info.ModTime(), info.Size()
	}
	prefix := Prefix + "/" + kindEngine + "/" + engine + "/" + version + "/"
	out := map[string]string{}
	for p, sum := range t.sums {
		if rest, ok := strings.CutPrefix(p, prefix); ok && path.Base(rest) == rest {
			out[rest] = sum
		}
	}
	return out
}

// Err is why the file could not be read, the last time it was.
func (t *Trusted) Err() error {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.err
}