llmgw/internal/bundle/read.go

package bundle

import (
	"archive/tar"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// maxManifest bounds the manifest read into memory.
const maxManifest = 16 << 20

// source is a bundle being read: a directory, or the same packed as a tar
// with the manifest first, read as a stream so tens of gigabytes of weights
// are never unpacked twice.
type source interface {
	// manifest is the manifest's bytes, as they are in the bundle.
	manifest() []byte
	// each hands over every file, in the order the bundle has them. A file
	// fn does not read is passed over.
	each(fn func(path string, r io.Reader) error) error
	close() error
}

// open reads a bundle's manifest, from a directory or a tar.
func open(p string) (source, Manifest, error) {
	info, err := os.Stat(p)
	if err != nil {
		return nil, Manifest{}, fmt.Errorf("the bundle: %w", err)
	}
	var src source
	if info.IsDir() {
		src, err = openDir(p)
	} else {
		src, err = openTar(p)
	}
	if err != nil {
		return nil, Manifest{}, err
	}
	m, err := parseManifest(src.manifest())
	if err != nil {
		_ = src.close()
		return nil, Manifest{}, err
	}
	if d, ok := src.(*dirSource); ok {
		d.files = m.Files
	}
	return src, m, nil
}

// ReadManifest reads a bundle's manifest and its digest, without reading
// its files.
func ReadManifest(p string) (Manifest, string, error) {
	src, m, err := open(p)
	if err != nil {
		return Manifest{}, "", err
	}
	defer src.close() //nolint:errcheck // read only
	return m, Digest(src.manifest()), nil
}

type dirSource struct {
	dir   string
	raw   []byte
	files []File
}

func openDir(dir string) (*dirSource, error) {
	raw, err := readLimited(filepath.Join(dir, ManifestName))
	if err != nil {
		return nil, fmt.Errorf("%s is not a bundle: %w", dir, err)
	}
	return &dirSource{dir: dir, raw: raw}, nil
}

func (d *dirSource) manifest() []byte { return d.raw }
func (d *dirSource) close() error     { return nil }

func (d *dirSource) each(fn func(string, io.Reader) error) error {
	for _, f := range d.files {
		// The path was checked when the manifest was read: plain names
		// under the bundle's three directories.
		in, err := os.Open(filepath.Join(d.dir, filepath.FromSlash(f.Path)))
		if err != nil {
			return fmt.Errorf("the bundle lists %s and does not have it: %w", f.Path, err)
		}
		err = fn(f.Path, in)
		in.Close()
		if err != nil {
			return err
		}
	}
	return nil
}

type tarSource struct {
	f   *os.File
	tr  *tar.Reader
	raw []byte
}

func openTar(p string) (*tarSource, error) {
	f, err := os.Open(p) //nolint:gosec // the operator's bundle
	if err != nil {
		return nil, err
	}
	tr := tar.NewReader(f)
	h, err := tr.Next()
	if err != nil || h.Name != ManifestName || h.Typeflag != tar.TypeReg {
		f.Close()
		return nil, fmt.Errorf("%s is not a bundle: its first entry is not %s", p, ManifestName)
	}
	raw, err := io.ReadAll(io.LimitReader(tr, maxManifest))
	if err != nil {
		f.Close()
		return nil, err
	}
	return &tarSource{f: f, tr: tr, raw: raw}, nil
}

func (t *tarSource) manifest() []byte { return t.raw }
func (t *tarSource) close() error     { return t.f.Close() }

func (t *tarSource) each(fn func(string, io.Reader) error) error {
	for {
		h, err := t.tr.Next()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return fmt.Errorf("reading the bundle: %w", err)
		}
		if h.Typeflag == tar.TypeDir {
			continue
		}
		if h.Typeflag != tar.TypeReg {
			// Links and devices have no place in a bundle; the files it
			// lists are all it holds.
			return fmt.Errorf("the bundle holds %s, which is not a plain file", h.Name)
		}
		if err := fn(strings.TrimPrefix(h.Name, "./"), t.tr); err != nil {
			return err
		}
	}
}

func readLimited(p string) ([]byte, error) {
	f, err := os.Open(p) //nolint:gosec // the operator's bundle
	if err != nil {
		return nil, err
	}
	defer f.Close()
	return io.ReadAll(io.LimitReader(f, maxManifest))
}