llmgw/internal/bundle/import.go

package bundle

import (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"maps"
	"os"
	"path"
	"path/filepath"
	"slices"
	"strings"

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

// ImportOptions say where a bundle goes and how it is known to be the one
// that was made.
type ImportOptions struct {
	// Mirror is the mirror's root directory, /srv/models. Empty checks the
	// bundle and writes nothing.
	Mirror string
	// SHA256 is the manifest's digest as `bundle create` printed it.
	// Unchecked takes the bundle without one, said by whoever runs it.
	SHA256    string
	Unchecked bool
	// Progress hears each file as it is done: kept is one the mirror
	// already had, the same.
	Progress func(f File, kept bool)
}

// Result is what an import did.
type Result struct {
	Manifest Manifest
	// Digest is the manifest's sha256.
	Digest string
	// Written and Kept count files: put in the mirror, and already there.
	Written, Kept int
	Bytes         int64
}

// partSuffix marks a file still being written; the mirror lists none.
const partSuffix = ".part"

// Import checks every file of a bundle and puts it in a mirror.
//
// Nothing is in place until everything is checked. Each file is written
// beside where it goes, under a name the mirror does not serve, and checked
// against the manifest as it is written; each release's checksums are
// checked against the project's signing key and its binaries against them.
// Only then is each renamed into place — so a bundle damaged anywhere
// leaves the mirror as it was.
//
// A file the mirror already has is kept when it is the same, and refused
// when it is not: a version in the mirror is never replaced, because that
// is what makes going back to it possible.
func Import(p string, o ImportOptions) (Result, error) {
	src, m, err := open(p)
	if err != nil {
		return Result{}, err
	}
	defer src.close() //nolint:errcheck // read only
	res := Result{Manifest: m, Digest: Digest(src.manifest())}
	if err := checkDigest(src.manifest(), o.SHA256, o.Unchecked); err != nil {
		return res, err
	}

	byPath := map[string]File{}
	dest := map[string]string{}
	keep := map[string]bool{}
	for _, f := range m.Files {
		byPath[f.Path] = f
		if o.Mirror == "" {
			continue
		}
		rel, _ := Destination(f.Path) // checked when the manifest was read
		dest[f.Path] = filepath.Join(o.Mirror, filepath.FromSlash(rel))
		same, err := alreadyHas(dest[f.Path], f)
		if err != nil {
			return res, err
		}
		keep[f.Path] = same
	}

	// Each release's checksums are kept as they pass, to be checked against
	// the signing key once the whole release has been read.
	small := map[string][]byte{}
	var written []string
	defer func() {
		// Whatever was staged and not put in place goes.
		for _, p := range written {
			_ = os.Remove(dest[p] + partSuffix)
		}
	}()
	seen := map[string]bool{}
	err = src.each(func(name string, r io.Reader) error {
		f, ok := byPath[name]
		if !ok {
			return fmt.Errorf("the bundle holds %s, which its manifest does not list", name)
		}
		if seen[name] {
			return fmt.Errorf("the bundle holds %s twice", name)
		}
		seen[name] = true
		var w io.Writer = io.Discard
		var out *os.File
		if o.Mirror != "" && !keep[name] {
			if err := mkdirsLike(o.Mirror, filepath.Dir(dest[name])); err != nil {
				return err
			}
			out, err = os.OpenFile(dest[name]+partSuffix, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) //nolint:gosec // a checked path under the mirror
			if err != nil {
				return err
			}
			written = append(written, name)
			w = out
		}
		if keep[name] {
			r = eofReader{}
		}
		var buf bytes.Buffer
		if isReleaseSums(name) {
			w = io.MultiWriter(w, &limited{w: &buf, n: 1 << 20})
		}
		h := sha256.New()
		n, err := io.Copy(io.MultiWriter(w, h), r)
		if out != nil {
			if cerr := out.Close(); err == nil {
				err = cerr
			}
		}
		if err != nil {
			return fmt.Errorf("%s: %w", name, err)
		}
		if keep[name] {
			if isReleaseSums(name) {
				b, err := os.ReadFile(dest[name])
				if err != nil {
					return err
				}
				small[name] = b
			}
			res.Kept++
			if o.Progress != nil {
				o.Progress(f, true)
			}
			return nil
		}
		if got := hex.EncodeToString(h.Sum(nil)); got != f.SHA256 || n != f.Bytes {
			return fmt.Errorf("%s is not what the bundle was made with (sha256 %s, expected %s): "+
				"the bundle was damaged on the way; nothing was put in the mirror", name, got[:12], f.SHA256[:12])
		}
		if isReleaseSums(name) {
			small[name] = buf.Bytes()
		}
		res.Written++
		res.Bytes += n
		if o.Progress != nil {
			o.Progress(f, false)
		}
		return nil
	})
	if err != nil {
		return res, err
	}
	for _, f := range m.Files {
		if !seen[f.Path] {
			return res, fmt.Errorf("the bundle lacks %s, which its manifest lists", f.Path)
		}
	}
	if err := checkReleases(m, small); err != nil {
		return res, err
	}
	if o.Mirror == "" {
		return res, nil
	}

	for _, name := range written {
		d := dest[name]
		f := byPath[name]
		if strings.HasPrefix(name, kindModel+"/") {
			// Beside the file before it is in place, as a download into the
			// mirror does, so a node is never offered it without its checksum.
			if err := mirror.Record(d, f.SHA256, f.Source); err != nil {
				return res, err
			}
		}
		if err := os.Rename(d+partSuffix, d); err != nil {
			return res, err
		}
		chownLike(o.Mirror, d)
		if strings.HasPrefix(name, kindModel+"/") {
			dir, base := filepath.Split(d)
			chownLike(o.Mirror, filepath.Join(dir, "."+base+".sha256"))
			chownLike(o.Mirror, filepath.Join(dir, "."+base+".source"))
		}
	}
	written = nil
	return res, nil
}

// alreadyHas says whether the mirror has this file already, the same; an
// error when it has a different one there.
func alreadyHas(dst string, f File) (bool, error) {
	info, err := os.Stat(dst)
	if errors.Is(err, os.ErrNotExist) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	sum := mirror.Recorded(dst)
	if sum == "" || info.Size() != f.Bytes {
		if sum, err = fileSHA256(dst); err != nil {
			return false, err
		}
	}
	if sum != f.SHA256 || info.Size() != f.Bytes {
		return false, fmt.Errorf("the mirror already has %s, and it is not the bundle's: "+
			"a version in the mirror is never replaced; give the model another version in a new bundle", dst)
	}
	return true, nil
}

func isReleaseSums(name string) bool {
	if !strings.HasPrefix(name, kindRelease+"/") {
		return false
	}
	b := path.Base(name)
	return b == "checksums.txt" || b == release.SignatureName
}

// checkReleases holds each release's binaries to its checksums, and the
// checksums to the project's signing key — as `node add` and an upgrade
// would online. The manifest's own checksums say the bundle arrived as it
// was made; these say what was made is the project's release.
func checkReleases(m Manifest, small map[string][]byte) error {
	byVersion := map[string][]File{}
	for _, f := range m.Files {
		if strings.HasPrefix(f.Path, kindRelease+"/") {
			v := strings.Split(f.Path, "/")[1]
			byVersion[v] = append(byVersion[v], f)
		}
	}
	for _, v := range slices.Sorted(maps.Keys(byVersion)) {
		root := kindRelease + "/" + v + "/"
		text, ok := small[root+"checksums.txt"]
		if !ok {
			return fmt.Errorf("the bundle's llmgw %s has no checksums.txt", v)
		}
		c := release.Checksums{Text: text, Signature: string(small[root+release.SignatureName])}
		if err := c.Verify(); err != nil {
			return fmt.Errorf("llmgw %s in the bundle: %w", v, err)
		}
		for _, f := range byVersion[v] {
			name := path.Base(f.Path)
			if name == "checksums.txt" || name == release.SignatureName {
				continue
			}
			want, err := c.For(name)
			if err != nil {
				return fmt.Errorf("llmgw %s in the bundle: %w", v, err)
			}
			if want != f.SHA256 {
				return fmt.Errorf("%s in the bundle is not the release's (the release gives sha256 %s)", f.Path, want[:12])
			}
		}
	}
	return nil
}

// mkdirsLike makes dir and every directory above it up to root, each owned
// as root is: the mirror runs as its own user, and removes what it holds.
func mkdirsLike(root, dir string) error {
	rel, err := filepath.Rel(root, dir)
	if err != nil || strings.HasPrefix(rel, "..") {
		return fmt.Errorf("%s is outside the mirror at %s", dir, root)
	}
	cur := root
	for _, e := range strings.Split(rel, string(filepath.Separator)) {
		if e == "." {
			continue
		}
		cur = filepath.Join(cur, e)
		if _, err := os.Stat(cur); err == nil {
			continue
		}
		if err := os.Mkdir(cur, 0o755); err != nil && !errors.Is(err, os.ErrExist) {
			return err
		}
		chownLike(root, cur)
	}
	return nil
}

func fileSHA256(p string) (string, error) {
	f, err := os.Open(p) //nolint:gosec // a file in the mirror
	if err != nil {
		return "", err
	}
	defer f.Close()
	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return "", err
	}
	return hex.EncodeToString(h.Sum(nil)), nil
}

// eofReader is a file passed over: a tar moves past it on its own.
type eofReader struct{}

func (eofReader) Read([]byte) (int, error) { return 0, io.EOF }

// limited keeps at most n bytes and quietly drops the rest: a checksums
// file larger than that is refused by its signature, not by a write error.
type limited struct {
	w io.Writer
	n int64
}

func (l *limited) Write(p []byte) (int, error) {
	if l.n <= 0 {
		return len(p), nil
	}
	k := min(int64(len(p)), l.n)
	l.n -= k
	if _, err := l.w.Write(p[:k]); err != nil {
		return 0, err
	}
	return len(p), nil
}