source/pkg/compactonnx/safe_files.go

package compactonnx

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

const maxMetadataBytes = 64 * 1024

type artifactPair struct {
	model          []byte
	metadata       []byte
	modelSHA256    string
	metadataSHA256 string
}

func readArtifactPair(modelPath, metadataPath string) (artifactPair, error) {
	if modelPath == "" {
		return artifactPair{}, errors.New("compact ONNX model is required; fallback is unavailable")
	}
	if hasTraversal(modelPath) || hasTraversal(metadataPath) {
		return artifactPair{}, errors.New("compact ONNX artifact path contains traversal")
	}
	modelAbs, err := filepath.Abs(modelPath)
	if err != nil {
		return artifactPair{}, fmt.Errorf("resolve compact ONNX path: %w", err)
	}
	if metadataPath == "" {
		metadataPath = modelPath + ".json"
	}
	metadataAbs, err := filepath.Abs(metadataPath)
	if err != nil {
		return artifactPair{}, fmt.Errorf("resolve compact metadata path: %w", err)
	}

	modelParent, err := filepath.EvalSymlinks(filepath.Dir(modelAbs))
	if err != nil {
		return artifactPair{}, fmt.Errorf("resolve compact ONNX directory: %w", err)
	}
	metadataParent, err := filepath.EvalSymlinks(filepath.Dir(metadataAbs))
	if err != nil {
		return artifactPair{}, fmt.Errorf("resolve compact metadata directory: %w", err)
	}
	if modelParent != metadataParent {
		return artifactPair{}, errors.New("compact metadata must be adjacent to the ONNX model")
	}
	modelName := filepath.Base(modelAbs)
	metadataName := filepath.Base(metadataAbs)
	if metadataName != modelName+".json" {
		return artifactPair{}, errors.New("compact metadata must use the adjacent <model>.onnx.json name")
	}
	if modelName == "." || modelName == string(filepath.Separator) || modelName == "" {
		return artifactPair{}, errors.New("invalid compact ONNX filename")
	}

	root, err := os.OpenRoot(modelParent)
	if err != nil {
		return artifactPair{}, fmt.Errorf("open compact artifact directory: %w", err)
	}
	defer root.Close()

	model, modelInfo, err := readRootFile(root, modelName, MaxModelBytes)
	if err != nil {
		return artifactPair{}, fmt.Errorf("read compact ONNX: %w", err)
	}
	metadata, _, err := readRootFile(root, metadataName, maxMetadataBytes)
	if err != nil {
		return artifactPair{}, fmt.Errorf("read compact metadata: %w", err)
	}
	if modelInfo.Size() == 0 {
		return artifactPair{}, errors.New("compact ONNX is empty")
	}
	modelHash := sha256.Sum256(model)
	metadataHash := sha256.Sum256(metadata)
	return artifactPair{
		model: model, metadata: metadata,
		modelSHA256:    hex.EncodeToString(modelHash[:]),
		metadataSHA256: hex.EncodeToString(metadataHash[:]),
	}, nil
}

func readRootFile(root *os.Root, name string, maximum int64) ([]byte, os.FileInfo, error) {
	if name != filepath.Base(name) || name == "." || name == ".." {
		return nil, nil, errors.New("artifact name is not a filename")
	}
	before, err := root.Lstat(name)
	if err != nil {
		return nil, nil, err
	}
	if before.Mode()&os.ModeSymlink != 0 || !before.Mode().IsRegular() {
		return nil, nil, errors.New("artifact is not an unlinked regular file")
	}
	if before.Size() < 0 || before.Size() > maximum {
		return nil, nil, fmt.Errorf("artifact size %d exceeds limit %d", before.Size(), maximum)
	}

	file, err := root.Open(name)
	if err != nil {
		return nil, nil, err
	}
	defer file.Close()
	opened, err := file.Stat()
	if err != nil {
		return nil, nil, err
	}
	if !opened.Mode().IsRegular() || !os.SameFile(before, opened) {
		return nil, nil, errors.New("artifact changed while opening")
	}
	if links, err := openedHardLinkCount(file, opened); err != nil || links != 1 {
		return nil, nil, errors.New("opened artifact link count is not exactly one")
	}
	data, err := io.ReadAll(io.LimitReader(file, maximum+1))
	if err != nil {
		return nil, nil, err
	}
	if int64(len(data)) != before.Size() {
		return nil, nil, errors.New("artifact size changed while reading")
	}
	final, err := file.Stat()
	if err != nil {
		return nil, nil, err
	}
	if final.Size() != before.Size() || !os.SameFile(opened, final) {
		return nil, nil, errors.New("artifact changed while reading")
	}
	if links, err := openedHardLinkCount(file, final); err != nil || links != 1 {
		return nil, nil, errors.New("final artifact link count is not exactly one")
	}
	return data, final, nil
}

func hasTraversal(value string) bool {
	if value == "" {
		return false
	}
	value = strings.ReplaceAll(value, "\\", "/")
	for _, part := range strings.Split(value, "/") {
		if part == ".." {
			return true
		}
	}
	return false
}