notarius/internal/framework/chunkplan/store.go

package chunkplan

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const SchemaVersion = pipeline.ChunkPlanSchemaVersion

const planFileName = "plan.json"

type filesystemStore struct {
	root  string
	write func(*os.Root, string, []byte) error
}

func NewFilesystemStore(root string) (pipeline.ChunkPlanStore, error) {
	root = strings.TrimSpace(root)
	if root == "" {
		return nil, fmt.Errorf("chunk plan root must not be empty")
	}
	if strings.ContainsRune(root, '\x00') {
		return nil, fmt.Errorf("chunk plan root must not contain NUL")
	}
	return &filesystemStore{root: filepath.Clean(root), write: writeAtomic}, nil
}

func (s *filesystemStore) Load(sourceDigest string) (pipeline.ChunkPlanRecord, pipeline.ChunkPlanDecision, error) {
	digestDir, err := digestPathSegment(sourceDigest)
	if err != nil {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, err
	}
	root, err := s.openRoot(false)
	if err != nil {
		if os.IsNotExist(err) {
			return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing, Reason: lookupReason(pipeline.ChunkPlanMissing)}, nil
		}
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, fmt.Errorf("open chunk plan root: %w", err)
	}
	defer root.Close()

	state, err := inspectDirectory(root, digestDir)
	if err != nil {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, fmt.Errorf("inspect chunk plan directory: %w", err)
	}
	if state == entryMissing {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing, Reason: lookupReason(pipeline.ChunkPlanMissing)}, nil
	}
	if state == entryRejected {
		return invalidDecision()
	}

	target := planPath(digestDir)
	state, err = inspectPlan(root, target)
	if err != nil {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, fmt.Errorf("inspect chunk plan file: %w", err)
	}
	if state == entryMissing {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing, Reason: lookupReason(pipeline.ChunkPlanMissing)}, nil
	}
	if state == entryRejected {
		return invalidDecision()
	}

	data, err := root.ReadFile(target)
	if err != nil {
		if os.IsNotExist(err) {
			return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing, Reason: lookupReason(pipeline.ChunkPlanMissing)}, nil
		}
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, fmt.Errorf("read chunk plan: %w", err)
	}

	var record pipeline.ChunkPlanRecord
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	decoder.UseNumber()
	if err := decoder.Decode(&record); err != nil {
		return invalidDecision()
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		return invalidDecision()
	}
	if err := validateRecord(record, sourceDigest); err != nil {
		return invalidDecision()
	}
	return record, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanHit, Reason: lookupReason(pipeline.ChunkPlanHit)}, nil
}

func (s *filesystemStore) Save(record pipeline.ChunkPlanRecord) error {
	if err := validateRecord(record, record.SourceDigest); err != nil {
		return fmt.Errorf("validate chunk plan record: %w", err)
	}
	digestDir, err := digestPathSegment(record.SourceDigest)
	if err != nil {
		return err
	}
	data, err := json.Marshal(record)
	if err != nil {
		return fmt.Errorf("encode chunk plan record: %w", err)
	}
	data = append(data, '\n')
	root, err := s.openRoot(true)
	if err != nil {
		return fmt.Errorf("open chunk plan root: %w", err)
	}
	defer root.Close()
	if err := ensureDirectory(root, digestDir); err != nil {
		return fmt.Errorf("prepare chunk plan directory: %w", err)
	}
	target := planPath(digestDir)
	state, err := inspectPlan(root, target)
	if err != nil {
		return fmt.Errorf("inspect chunk plan file: %w", err)
	}
	if state == entryRejected {
		return fmt.Errorf("chunk plan file has an unsupported type")
	}
	writer := s.write
	if writer == nil {
		writer = writeAtomic
	}
	if err := writer(root, target, data); err != nil {
		return fmt.Errorf("write chunk plan: %w", err)
	}
	return nil
}

func (s *filesystemStore) openRoot(create bool) (*os.Root, error) {
	if s == nil || strings.TrimSpace(s.root) == "" {
		return nil, fmt.Errorf("chunk plan store must not be nil")
	}
	if create {
		if err := os.MkdirAll(s.root, 0o700); err != nil {
			return nil, err
		}
	}
	root, err := os.OpenRoot(s.root)
	if err != nil {
		return nil, err
	}
	if !create {
		return root, nil
	}
	rootDirectory, err := root.Open(".")
	if err != nil {
		_ = root.Close()
		return nil, err
	}
	defer rootDirectory.Close()
	if err := rootDirectory.Chmod(0o700); err != nil {
		_ = root.Close()
		return nil, err
	}
	return root, nil
}

type entryState uint8

const (
	entryPresent entryState = iota
	entryMissing
	entryRejected
)

func inspectDirectory(root *os.Root, digestDir string) (entryState, error) {
	info, err := root.Lstat(digestDir)
	if err != nil {
		if os.IsNotExist(err) {
			return entryMissing, nil
		}
		return entryPresent, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
		return entryRejected, nil
	}
	return entryPresent, nil
}

func ensureDirectory(root *os.Root, digestDir string) error {
	for {
		state, err := inspectDirectory(root, digestDir)
		if err != nil {
			return err
		}
		switch state {
		case entryRejected:
			return fmt.Errorf("chunk plan directory has an unsupported type")
		case entryMissing:
			if err := root.Mkdir(digestDir, 0o700); err != nil && !errors.Is(err, os.ErrExist) {
				return err
			}
			continue
		}

		directory, err := root.Open(digestDir)
		if err != nil {
			return err
		}
		info, statErr := directory.Stat()
		if statErr == nil && !info.IsDir() {
			statErr = fmt.Errorf("chunk plan directory has an unsupported type")
		}
		if statErr == nil {
			statErr = directory.Chmod(0o700)
		}
		closeErr := directory.Close()
		if statErr != nil {
			return statErr
		}
		return closeErr
	}
}

func inspectPlan(root *os.Root, target string) (entryState, error) {
	info, err := root.Lstat(target)
	if err != nil {
		if os.IsNotExist(err) {
			return entryMissing, nil
		}
		return entryPresent, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return entryRejected, nil
	}
	return entryPresent, nil
}

func planPath(digestDir string) string {
	return digestDir + "/" + planFileName
}

func digestPathSegment(digest string) (string, error) {
	if !strings.HasPrefix(digest, "sha256:") {
		return "", fmt.Errorf("source digest must use sha256:<64 lowercase hex> format")
	}
	hexDigest := strings.TrimPrefix(digest, "sha256:")
	if len(hexDigest) != 64 || strings.ToLower(hexDigest) != hexDigest {
		return "", fmt.Errorf("source digest must use sha256:<64 lowercase hex> format")
	}
	decoded, err := hex.DecodeString(hexDigest)
	if err != nil || len(decoded) != 32 {
		return "", fmt.Errorf("source digest must use sha256:<64 lowercase hex> format")
	}
	return hexDigest, nil
}

func validateRecord(record pipeline.ChunkPlanRecord, requestedDigest string) error {
	if record.SchemaVersion != SchemaVersion {
		return fmt.Errorf("schema_version %q is not supported", record.SchemaVersion)
	}
	if _, err := digestPathSegment(requestedDigest); err != nil {
		return err
	}
	if record.SourceDigest != requestedDigest {
		return fmt.Errorf("source_digest does not match requested source")
	}
	if record.Plan.SourceDigest != record.SourceDigest {
		return fmt.Errorf("plan source_digest does not match record source_digest")
	}
	if len(record.Plan.Ranges) == 0 {
		return fmt.Errorf("plan ranges must not be empty")
	}
	if err := source.ValidateChunkAnnotations(record.Plan.Annotations); err != nil {
		return fmt.Errorf("plan annotations: %w", err)
	}
	seenStarts := make(map[int]struct{}, len(record.Plan.Ranges))
	for i, chunkRange := range record.Plan.Ranges {
		if chunkRange.StartUnitID <= 0 || chunkRange.EndUnitID <= 0 {
			return fmt.Errorf("plan range[%d] boundaries must be positive", i)
		}
		if _, exists := seenStarts[chunkRange.StartUnitID]; exists {
			return fmt.Errorf("plan range[%d] duplicates start_unit_id %d", i, chunkRange.StartUnitID)
		}
		seenStarts[chunkRange.StartUnitID] = struct{}{}
		if err := source.ValidateChunkAnnotations(chunkRange.Annotations); err != nil {
			return fmt.Errorf("plan range[%d] annotations: %w", i, err)
		}
	}
	wantDigest, err := source.DigestChunkPlan(record.Plan)
	if err != nil {
		return fmt.Errorf("digest plan: %w", err)
	}
	if record.PlanDigest != wantDigest {
		return fmt.Errorf("plan_digest does not match plan")
	}
	if strings.TrimSpace(record.Producer.InputModule) == "" {
		return fmt.Errorf("producer input_module must not be empty")
	}
	if strings.TrimSpace(record.Producer.ChunkModule) == "" {
		return fmt.Errorf("producer chunk_module must not be empty")
	}
	if record.CreatedAt.IsZero() {
		return fmt.Errorf("created_at must not be zero")
	}
	return nil
}

func invalidDecision() (pipeline.ChunkPlanRecord, pipeline.ChunkPlanDecision, error) {
	return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanInvalid, Reason: lookupReason(pipeline.ChunkPlanInvalid)}, nil
}

func lookupReason(status pipeline.ChunkPlanStatus) string {
	switch status {
	case pipeline.ChunkPlanHit:
		return "stored chunk plan is valid"
	case pipeline.ChunkPlanMissing:
		return "chunk plan not found"
	case pipeline.ChunkPlanInvalid:
		return "stored chunk plan is invalid"
	default:
		return "chunk plan lookup skipped"
	}
}

type atomicWriteHooks struct {
	BeforeCreateTemp func() error
	BeforeRename     func() error
}

func writeAtomic(root *os.Root, target string, data []byte) error {
	return writeAtomicWithHooks(root, target, data, atomicWriteHooks{})
}

func writeAtomicWithHooks(root *os.Root, target string, data []byte, hooks atomicWriteHooks) error {
	if hooks.BeforeCreateTemp != nil {
		if err := hooks.BeforeCreateTemp(); err != nil {
			return err
		}
	}
	temp, tempPath, err := createTemporaryFile(root, target)
	if err != nil {
		return err
	}
	removeTemp := true
	defer func() {
		if removeTemp {
			_ = root.Remove(tempPath)
		}
	}()
	if err := temp.Chmod(0o600); err != nil {
		_ = temp.Close()
		return err
	}
	if _, err := temp.Write(data); err != nil {
		_ = temp.Close()
		return err
	}
	if err := temp.Sync(); err != nil {
		_ = temp.Close()
		return err
	}
	if err := temp.Close(); err != nil {
		return err
	}
	if hooks.BeforeRename != nil {
		if err := hooks.BeforeRename(); err != nil {
			return err
		}
	}
	if err := root.Rename(tempPath, target); err != nil {
		return err
	}
	removeTemp = false
	return nil
}

func createTemporaryFile(root *os.Root, target string) (*os.File, string, error) {
	for attempt := 0; attempt < 32; attempt++ {
		var suffix [16]byte
		if _, err := rand.Read(suffix[:]); err != nil {
			return nil, "", err
		}
		path := strings.TrimSuffix(target, planFileName) + ".plan.json.tmp-" + hex.EncodeToString(suffix[:])
		file, err := root.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
		if errors.Is(err, os.ErrExist) {
			continue
		}
		if err != nil {
			return nil, "", err
		}
		return file, path, nil
	}
	return nil, "", fmt.Errorf("create unique temporary chunk plan file")
}