iop/packages/go/agentstate/store.go

// Package agentstate provides the crash-safe device-local persistence boundary
// for the shared AgentTask manager state.
package agentstate

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"

	"iop/packages/go/agenttask"
)

const storeSchemaVersion uint32 = 1
const maxIntegrationRecordPayload = 16 << 20

var (
	ErrCorruptState      = errors.New("agentstate corrupt state")
	ErrUnsupportedSchema = errors.New("agentstate unsupported store schema")
)

type StateError struct {
	Path string
	Kind error
	Err  error
}

func (e *StateError) Error() string {
	return fmt.Sprintf("agentstate: %v at %q: %v", e.Kind, e.Path, e.Err)
}

func (e *StateError) Unwrap() []error {
	return []error{e.Kind, e.Err}
}

type Store struct {
	path     string
	lockPath string
}

// IntegrationRecordSnapshot is an immutable view of one checksum-covered
// integration record and its key-local content revision.
type IntegrationRecordSnapshot struct {
	Payload  []byte
	Revision string
}

// IntegrationRecordUpdate describes one key-local compare-and-swap update.
// All updates supplied to CompareAndSwapIntegrationRecords commit together.
type IntegrationRecordUpdate struct {
	Key      string
	Expected string
	Payload  []byte
}

func NewStore(path string) (*Store, error) {
	if path == "" {
		return nil, fmt.Errorf("agentstate: state path is required")
	}
	clean := filepath.Clean(path)
	if clean == "." {
		return nil, fmt.Errorf("agentstate: state path must name a file")
	}
	return &Store{path: clean, lockPath: clean + ".lock"}, nil
}

func (s *Store) Path() string {
	return s.path
}

// LoadIntegrationRecord returns one opaque, checksum-covered host integration
// journal together with its key-local content revision.
func (s *Store) LoadIntegrationRecord(
	ctx context.Context,
	key string,
) ([]byte, string, bool, error) {
	if err := validateIntegrationKey(key); err != nil {
		return nil, "", false, err
	}
	if err := ctx.Err(); err != nil {
		return nil, "", false, err
	}
	var payload []byte
	var found bool
	err := s.withLock(ctx, syscall.LOCK_SH, func() error {
		_, _, records, err := s.readUnlocked()
		if err != nil {
			return err
		}
		record, ok := records[key]
		if !ok {
			return nil
		}
		payload = append([]byte(nil), record...)
		found = true
		return nil
	})
	if err != nil {
		return nil, "", false, err
	}
	if !found {
		return nil, "", false, nil
	}
	return payload, integrationRecordRevision(payload), true, nil
}

// LoadIntegrationRecords returns cloned, checksum-covered integration records
// whose keys start with prefix.
func (s *Store) LoadIntegrationRecords(
	ctx context.Context,
	prefix string,
) (map[string]IntegrationRecordSnapshot, error) {
	if err := validateIntegrationKey(prefix); err != nil {
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	snapshots := make(map[string]IntegrationRecordSnapshot)
	err := s.withLock(ctx, syscall.LOCK_SH, func() error {
		_, _, records, err := s.readUnlocked()
		if err != nil {
			return err
		}
		for key, record := range records {
			if !strings.HasPrefix(key, prefix) {
				continue
			}
			payload := append([]byte(nil), record...)
			snapshots[key] = IntegrationRecordSnapshot{
				Payload:  payload,
				Revision: integrationRecordRevision(payload),
			}
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return snapshots, nil
}

// CompareAndSwapIntegrationRecord atomically updates one opaque integration
// journal while preserving the manager snapshot and every sibling journal.
func (s *Store) CompareAndSwapIntegrationRecord(
	ctx context.Context,
	key string,
	expected string,
	payload []byte,
) (string, error) {
	revisions, err := s.CompareAndSwapIntegrationRecords(ctx, []IntegrationRecordUpdate{{
		Key:      key,
		Expected: expected,
		Payload:  payload,
	}})
	if err != nil {
		return "", err
	}
	return revisions[key], nil
}

// CompareAndSwapIntegrationRecords atomically updates multiple opaque
// integration records while preserving the manager snapshot and all siblings.
// Every expected revision is checked before any record is mutated.
func (s *Store) CompareAndSwapIntegrationRecords(
	ctx context.Context,
	updates []IntegrationRecordUpdate,
) (map[string]string, error) {
	if len(updates) == 0 {
		return nil, fmt.Errorf("agentstate: at least one integration record update is required")
	}
	prepared := make([]IntegrationRecordUpdate, len(updates))
	seenKeys := make(map[string]struct{}, len(updates))
	for index, update := range updates {
		if err := validateIntegrationKey(update.Key); err != nil {
			return nil, fmt.Errorf("agentstate: integration record update %d: %w", index, err)
		}
		if _, duplicate := seenKeys[update.Key]; duplicate {
			return nil, fmt.Errorf("agentstate: duplicate integration record key %q", update.Key)
		}
		seenKeys[update.Key] = struct{}{}
		if len(update.Payload) == 0 ||
			len(update.Payload) > maxIntegrationRecordPayload ||
			!json.Valid(update.Payload) {
			return nil, fmt.Errorf(
				"agentstate: integration record update %d must contain one bounded JSON value",
				index,
			)
		}
		var decoded any
		if err := decodeOne(update.Payload, &decoded); err != nil {
			return nil, fmt.Errorf(
				"agentstate: decode integration record update %d: %w",
				index,
				err,
			)
		}
		prepared[index] = IntegrationRecordUpdate{
			Key:      update.Key,
			Expected: update.Expected,
			Payload:  append([]byte(nil), update.Payload...),
		}
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	revisions := make(map[string]string, len(prepared))
	err := s.withLock(ctx, syscall.LOCK_EX, func() error {
		state, current, records, err := s.readUnlocked()
		if err != nil {
			return err
		}
		for _, update := range prepared {
			currentRecord, exists := records[update.Key]
			currentRevision := ""
			if exists {
				currentRevision = integrationRecordRevision(currentRecord)
			}
			if currentRevision != update.Expected {
				return agenttask.ErrRevisionConflict
			}
		}

		changed := false
		if records == nil {
			records = make(map[string]json.RawMessage)
		}
		for _, update := range prepared {
			revisions[update.Key] = integrationRecordRevision(update.Payload)
			if currentRecord, exists := records[update.Key]; exists &&
				bytes.Equal(currentRecord, update.Payload) {
				continue
			}
			records[update.Key] = append(json.RawMessage(nil), update.Payload...)
			changed = true
		}
		if !changed {
			return nil
		}
		if current == ^uint64(0) {
			return fmt.Errorf("agentstate: revision overflow")
		}
		return s.writeUnlocked(current+1, state, records)
	})
	if err != nil {
		return nil, err
	}
	return revisions, nil
}

func (s *Store) Load(
	ctx context.Context,
) (agenttask.ManagerState, agenttask.StateRevision, error) {
	if err := ctx.Err(); err != nil {
		return agenttask.ManagerState{}, "", err
	}
	var state agenttask.ManagerState
	var revision uint64
	err := s.withLock(ctx, syscall.LOCK_SH, func() error {
		var err error
		state, revision, _, err = s.readUnlocked()
		return err
	})
	if err != nil {
		return agenttask.ManagerState{}, "", err
	}
	return state, agenttask.StateRevision(strconv.FormatUint(revision, 10)), nil
}

func (s *Store) CompareAndSwap(
	ctx context.Context,
	expected agenttask.StateRevision,
	next agenttask.ManagerState,
) (agenttask.StateRevision, error) {
	if err := ctx.Err(); err != nil {
		return "", err
	}
	var committed uint64
	err := s.withLock(ctx, syscall.LOCK_EX, func() error {
		_, current, integrationRecords, err := s.readUnlocked()
		if err != nil {
			return err
		}
		currentRevision := agenttask.StateRevision(strconv.FormatUint(current, 10))
		if expected != currentRevision {
			return agenttask.ErrRevisionConflict
		}
		if current == ^uint64(0) {
			return fmt.Errorf("agentstate: revision overflow")
		}
		if next.SchemaVersion != agenttask.StateSchemaVersion {
			return s.stateError(
				ErrUnsupportedSchema,
				fmt.Errorf("manager schema version %d", next.SchemaVersion),
			)
		}
		committed = current + 1
		return s.writeUnlocked(committed, next, integrationRecords)
	})
	if err != nil {
		return "", err
	}
	return agenttask.StateRevision(strconv.FormatUint(committed, 10)), nil
}

type diskEnvelope struct {
	SchemaVersion      uint32                     `json:"schema_version"`
	Revision           uint64                     `json:"revision"`
	Checksum           string                     `json:"checksum"`
	State              json.RawMessage            `json:"state"`
	IntegrationRecords map[string]json.RawMessage `json:"integration_records,omitempty"`
}

type checksumPayload struct {
	SchemaVersion      uint32                     `json:"schema_version"`
	Revision           uint64                     `json:"revision"`
	State              json.RawMessage            `json:"state"`
	IntegrationRecords map[string]json.RawMessage `json:"integration_records,omitempty"`
}

func (s *Store) readUnlocked() (
	agenttask.ManagerState,
	uint64,
	map[string]json.RawMessage,
	error,
) {
	payload, err := os.ReadFile(s.path)
	if errors.Is(err, os.ErrNotExist) {
		return agenttask.ManagerState{SchemaVersion: agenttask.StateSchemaVersion},
			0,
			nil,
			nil
	}
	if err != nil {
		return agenttask.ManagerState{}, 0, nil, err
	}
	var envelope diskEnvelope
	if err := decodeOne(payload, &envelope); err != nil {
		return agenttask.ManagerState{}, 0, nil, s.stateError(
			ErrCorruptState,
			fmt.Errorf("decode envelope: %w", err),
		)
	}
	if envelope.SchemaVersion != storeSchemaVersion {
		return agenttask.ManagerState{}, 0, nil, s.stateError(
			ErrUnsupportedSchema,
			fmt.Errorf("schema version %d", envelope.SchemaVersion),
		)
	}
	expected, err := stateChecksum(
		envelope.SchemaVersion,
		envelope.Revision,
		envelope.State,
		envelope.IntegrationRecords,
	)
	if err != nil {
		return agenttask.ManagerState{}, 0, nil, s.stateError(ErrCorruptState, err)
	}
	if envelope.Checksum == "" || !bytes.Equal(
		[]byte(envelope.Checksum),
		[]byte(expected),
	) {
		return agenttask.ManagerState{}, 0, nil, s.stateError(
			ErrCorruptState,
			fmt.Errorf("checksum mismatch"),
		)
	}
	var state agenttask.ManagerState
	if err := decodeOne(envelope.State, &state); err != nil {
		return agenttask.ManagerState{}, 0, nil, s.stateError(
			ErrCorruptState,
			fmt.Errorf("decode manager state: %w", err),
		)
	}
	for key, record := range envelope.IntegrationRecords {
		if validateIntegrationKey(key) != nil || len(record) == 0 || !json.Valid(record) {
			return agenttask.ManagerState{}, 0, nil, s.stateError(
				ErrCorruptState,
				fmt.Errorf("invalid integration record %q", key),
			)
		}
	}
	return state, envelope.Revision, cloneRawRecords(envelope.IntegrationRecords), nil
}

func (s *Store) writeUnlocked(
	revision uint64,
	state agenttask.ManagerState,
	integrationRecords map[string]json.RawMessage,
) error {
	statePayload, err := json.Marshal(state)
	if err != nil {
		return fmt.Errorf("agentstate: encode manager state: %w", err)
	}
	checksum, err := stateChecksum(
		storeSchemaVersion,
		revision,
		statePayload,
		integrationRecords,
	)
	if err != nil {
		return err
	}
	envelopePayload, err := json.Marshal(diskEnvelope{
		SchemaVersion:      storeSchemaVersion,
		Revision:           revision,
		Checksum:           checksum,
		State:              statePayload,
		IntegrationRecords: integrationRecords,
	})
	if err != nil {
		return fmt.Errorf("agentstate: encode envelope: %w", err)
	}
	envelopePayload = append(envelopePayload, '\n')

	dir := filepath.Dir(s.path)
	temp, err := os.CreateTemp(dir, "."+filepath.Base(s.path)+".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(envelopePayload); 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, s.path); err != nil {
		return err
	}
	removeTemp = false
	dirHandle, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer dirHandle.Close()
	return dirHandle.Sync()
}

func (s *Store) withLock(
	ctx context.Context,
	mode int,
	action func() error,
) error {
	if err := os.MkdirAll(filepath.Dir(s.path), 0o700); err != nil {
		return err
	}
	lock, err := os.OpenFile(s.lockPath, os.O_CREATE|os.O_RDWR, 0o600)
	if err != nil {
		return err
	}
	defer lock.Close()
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := syscall.Flock(int(lock.Fd()), mode); err != nil {
		return err
	}
	defer syscall.Flock(int(lock.Fd()), syscall.LOCK_UN) //nolint:errcheck
	if err := ctx.Err(); err != nil {
		return err
	}
	return action()
}

func stateChecksum(
	schemaVersion uint32,
	revision uint64,
	state json.RawMessage,
	integrationRecords ...map[string]json.RawMessage,
) (string, error) {
	var records map[string]json.RawMessage
	if len(integrationRecords) != 0 {
		records = integrationRecords[0]
	}
	payload, err := json.Marshal(checksumPayload{
		SchemaVersion:      schemaVersion,
		Revision:           revision,
		State:              state,
		IntegrationRecords: records,
	})
	if err != nil {
		return "", fmt.Errorf("agentstate: encode checksum payload: %w", err)
	}
	sum := sha256.Sum256(payload)
	return hex.EncodeToString(sum[:]), nil
}

func decodeOne(payload []byte, target any) error {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return err
	}
	var trailing any
	if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) {
		if err == nil {
			return fmt.Errorf("multiple JSON values")
		}
		return err
	}
	return nil
}

func validateIntegrationKey(key string) error {
	if strings.TrimSpace(key) == "" ||
		strings.TrimSpace(key) != key ||
		strings.ContainsAny(key, "\x00\r\n") {
		return fmt.Errorf("agentstate: invalid integration record key")
	}
	return nil
}

func integrationRecordRevision(payload []byte) string {
	sum := sha256.Sum256(payload)
	return "sha256:" + hex.EncodeToString(sum[:])
}

func cloneRawRecords(
	records map[string]json.RawMessage,
) map[string]json.RawMessage {
	if records == nil {
		return nil
	}
	copy := make(map[string]json.RawMessage, len(records))
	for key, record := range records {
		copy[key] = append(json.RawMessage(nil), record...)
	}
	return copy
}

func (s *Store) stateError(kind error, err error) error {
	return &StateError{Path: s.path, Kind: kind, Err: err}
}