llmgw/internal/mirror/import.go

package mirror

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// Import copies a file on this machine into the mirror at model/version/file:
// weights somebody already has — a GGUF they converted, Ollama's blob of
// one — put where every node can fetch them, with the checksum nodes check
// them against.
//
// The same discipline as Add: written under a temporary name and renamed
// when whole, the checksum and the source beside it first, and an existing
// version never replaced. wantSHA256, when given, is what the file must hash
// to — Ollama names each blob by its digest — and a file that does not is
// refused rather than copied to every node.
//
// source is what is recorded as where the file came from.
func Import(ctx context.Context, root, src string, e Entry, wantSHA256, source string,
	progress func(done, total int64),
) (Added, error) {
	if e.Model == "" || e.Version == "" || e.File == "" {
		return Added{}, fmt.Errorf("a model, a version and a filename are all required")
	}
	for _, n := range []string{e.Model, e.Version, e.File} {
		if err := ValidName(n); err != nil {
			return Added{}, err
		}
	}
	in, err := os.Open(src) //nolint:gosec // a file the operator named
	if err != nil {
		return Added{}, err
	}
	defer in.Close()
	st, err := in.Stat()
	if err != nil {
		return Added{}, err
	}
	if !st.Mode().IsRegular() {
		return Added{}, fmt.Errorf("%s is not a file", src)
	}
	if strings.HasSuffix(strings.ToLower(e.File), ".gguf") {
		if err := isGGUF(in, src); err != nil {
			return Added{}, err
		}
	}

	final := e.Path(root)
	if _, err := os.Stat(final); err == nil {
		return Added{}, fmt.Errorf("%s/%s/%s: %w", e.Model, e.Version, e.File, ErrExists)
	}
	if err := os.MkdirAll(filepath.Dir(final), 0o755); err != nil { //nolint:gosec // served to every node
		return Added{}, err
	}

	part := final + partSuffix
	// Truncated rather than refused when there: a .part is what an import cut
	// short left, and the next attempt is the one that should finish it.
	out, err := os.OpenFile(part, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) //nolint:gosec // inside the mirror root
	if err != nil {
		return Added{}, err
	}
	// Hashed on the way through: a forty gigabyte file is read once, not
	// copied and then read again.
	h := sha256.New()
	written, err := copyWithProgress(io.MultiWriter(out, h), ctxReader{ctx, in}, st.Size(), progress)
	closeErr := out.Close()
	if err == nil {
		err = closeErr
	}
	if err == nil && written != st.Size() {
		err = fmt.Errorf("%s: copied %d bytes of %d; it changed while being read", src, written, st.Size())
	}
	sum := hex.EncodeToString(h.Sum(nil))
	if err == nil && wantSHA256 != "" && !strings.EqualFold(sum, wantSHA256) {
		err = fmt.Errorf("%s is damaged: it should hash to sha256 %s and hashes to %s; "+
			"nothing was added", src, wantSHA256, sum)
	}
	if err == nil {
		err = writeChecksum(final, sum)
	}
	if err == nil {
		err = writeSource(final, source)
	}
	if err == nil {
		err = os.Rename(part, final)
	}
	if err != nil {
		_ = os.Remove(part)
		return Added{}, err
	}
	e.Bytes = written
	return Added{Entry: e, SHA256: sum}, nil
}

// ggufMagic opens every GGUF file.
const ggufMagic = "GGUF"

// isGGUF checks a file named .gguf is one, and leaves it read from the start.
// A file that is not — an HTML error page saved under the name, Ollama's
// newer safetensors-based layers — would otherwise reach every node and fail
// there, in llama.cpp's words.
func isGGUF(f *os.File, name string) error {
	head := make([]byte, len(ggufMagic))
	if _, err := io.ReadFull(f, head); err != nil || string(head) != ggufMagic {
		return fmt.Errorf("%s is not a GGUF file: it does not begin with %q", name, ggufMagic)
	}
	_, err := f.Seek(0, io.SeekStart)
	return err
}

// ctxReader stops a copy when its context ends: an operator's ^C on a long
// import leaves no partial file behind.
type ctxReader struct {
	ctx context.Context
	r   io.Reader
}

func (c ctxReader) Read(p []byte) (int, error) {
	if err := c.ctx.Err(); err != nil {
		return 0, err
	}
	return c.r.Read(p)
}