notarius/internal/framework/chunkplan/store.go

package chunkplan

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

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

const SchemaVersion = "notarius.chunk-plan.v1"

type filesystemStore struct {
	root string
}

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)}, nil
}

func (s *filesystemStore) Load(sourceDigest string) (pipeline.ChunkPlanRecord, pipeline.ChunkPlanDecision, error) {
	target, err := s.planPath(sourceDigest)
	if err != nil {
		return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{}, err
	}
	data, err := os.ReadFile(target)
	if err != nil {
		if os.IsNotExist(err) {
			return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanMissing, Reason: "chunk plan not found"}, 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(fmt.Sprintf("decode stored chunk plan: %v", err))
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return invalidDecision("stored chunk plan contains trailing JSON")
		}
		return invalidDecision(fmt.Sprintf("decode stored chunk plan trailer: %v", err))
	}
	if err := validateRecord(record, sourceDigest); err != nil {
		return invalidDecision(err.Error())
	}
	return record, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanHit, Reason: "stored chunk plan is valid"}, 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)
	}
	target, err := s.planPath(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')
	if err := writeAtomic(target, data); err != nil {
		return fmt.Errorf("write chunk plan: %w", err)
	}
	return nil
}

func (s *filesystemStore) planPath(sourceDigest string) (string, error) {
	if s == nil || strings.TrimSpace(s.root) == "" {
		return "", fmt.Errorf("chunk plan store must not be nil")
	}
	hexDigest, err := digestPathSegment(sourceDigest)
	if err != nil {
		return "", err
	}
	return filepath.Join(s.root, hexDigest, "plan.json"), nil
}

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(reason string) (pipeline.ChunkPlanRecord, pipeline.ChunkPlanDecision, error) {
	return pipeline.ChunkPlanRecord{}, pipeline.ChunkPlanDecision{Status: pipeline.ChunkPlanInvalid, Reason: reason}, nil
}

func writeAtomic(target string, data []byte) error {
	dir := filepath.Dir(target)
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return err
	}
	if err := os.Chmod(dir, 0o700); err != nil {
		return err
	}

	temp, err := os.CreateTemp(dir, ".plan.json.tmp-*")
	if err != nil {
		return err
	}
	tempPath := temp.Name()
	removeTemp := true
	defer func() {
		if removeTemp {
			_ = os.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 err := os.Rename(tempPath, target); err != nil {
		return err
	}
	removeTemp = false
	return nil
}