browser/platform/service/artifact_body_store.go

package service

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"browser.local/platform/domain"
	"browser.local/platform/repo"
	"browser.local/platform/validator"
)

type ArtifactBodyStore interface {
	SaveTransfer(domain.ArtifactTransferSession) error
	LoadTransfers() ([]domain.ArtifactTransferSession, error)
	PutPayload(string, []byte) error
	GetPayload(string) ([]byte, error)
}

type MemoryArtifactBodyStore struct {
	mu        sync.Mutex
	transfers map[string]domain.ArtifactTransferSession
	payloads  map[string][]byte
}

func NewMemoryArtifactBodyStore() *MemoryArtifactBodyStore {
	return &MemoryArtifactBodyStore{transfers: map[string]domain.ArtifactTransferSession{}, payloads: map[string][]byte{}}
}

func (store *MemoryArtifactBodyStore) SaveTransfer(session domain.ArtifactTransferSession) error {
	store.mu.Lock()
	defer store.mu.Unlock()
	store.transfers[session.TransferID] = domain.CopyArtifactTransferSession(session)
	return nil
}

func (store *MemoryArtifactBodyStore) LoadTransfers() ([]domain.ArtifactTransferSession, error) {
	store.mu.Lock()
	defer store.mu.Unlock()
	ids := make([]string, 0, len(store.transfers))
	for id := range store.transfers {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	out := make([]domain.ArtifactTransferSession, 0, len(ids))
	for _, id := range ids {
		out = append(out, domain.CopyArtifactTransferSession(store.transfers[id]))
	}
	return out, nil
}

func (store *MemoryArtifactBodyStore) PutPayload(artifactID string, payload []byte) error {
	store.mu.Lock()
	defer store.mu.Unlock()
	store.payloads[artifactID] = domain.CopyBytes(payload)
	return nil
}

func (store *MemoryArtifactBodyStore) GetPayload(artifactID string) ([]byte, error) {
	store.mu.Lock()
	defer store.mu.Unlock()
	payload, exists := store.payloads[artifactID]
	if !exists {
		return nil, repo.ErrNotFound
	}
	return domain.CopyBytes(payload), nil
}

type FileArtifactBodyStore struct {
	mu      sync.Mutex
	rootDir string
}

func NewFileArtifactBodyStore(rootDir string) (*FileArtifactBodyStore, error) {
	rootDir = strings.TrimSpace(rootDir)
	if rootDir == "" {
		return nil, fmt.Errorf("artifact directory is required")
	}
	for _, path := range []string{rootDir, filepath.Join(rootDir, "transfers"), filepath.Join(rootDir, "payloads")} {
		if err := os.MkdirAll(path, 0o700); err != nil {
			return nil, fmt.Errorf("create artifact body directory: %w", err)
		}
	}
	return &FileArtifactBodyStore{rootDir: rootDir}, nil
}

func (store *FileArtifactBodyStore) SaveTransfer(session domain.ArtifactTransferSession) error {
	store.mu.Lock()
	defer store.mu.Unlock()

	dir := store.transferDir(session.TransferID)
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return fmt.Errorf("create artifact transfer directory: %w", err)
	}
	manifest := domain.CopyArtifactTransferSession(session)
	for index, record := range manifest.ReceivedChunks {
		payload := domain.CopyBytes(record.Payload)
		if len(payload) != record.SizeBytes || validator.BytesChecksum(payload) != record.Checksum {
			return validationError("artifact chunk does not match durable manifest")
		}
		if err := writeAtomicFile(filepath.Join(dir, fmt.Sprintf("chunk-%08d.bin", index)), payload, 0o600); err != nil {
			return err
		}
		record.Payload = nil
		manifest.ReceivedChunks[index] = record
	}
	body, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return fmt.Errorf("encode artifact transfer manifest: %w", err)
	}
	return writeAtomicFile(filepath.Join(dir, "manifest.json"), body, 0o600)
}

func (store *FileArtifactBodyStore) LoadTransfers() ([]domain.ArtifactTransferSession, error) {
	store.mu.Lock()
	defer store.mu.Unlock()

	entries, err := os.ReadDir(filepath.Join(store.rootDir, "transfers"))
	if err != nil {
		return nil, fmt.Errorf("read artifact transfer directory: %w", err)
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].Name() < entries[j].Name() })
	out := make([]domain.ArtifactTransferSession, 0, len(entries))
	for _, entry := range entries {
		if !entry.IsDir() {
			continue
		}
		dir := filepath.Join(store.rootDir, "transfers", entry.Name())
		body, err := os.ReadFile(filepath.Join(dir, "manifest.json"))
		if err != nil {
			if errors.Is(err, os.ErrNotExist) {
				continue
			}
			return nil, fmt.Errorf("read artifact transfer manifest: %w", err)
		}
		var session domain.ArtifactTransferSession
		if err := json.Unmarshal(body, &session); err != nil {
			return nil, fmt.Errorf("decode artifact transfer manifest: %w", err)
		}
		if session.TransferID == "" || store.transferDir(session.TransferID) != dir {
			return nil, fmt.Errorf("artifact transfer manifest identity mismatch")
		}
		for index, record := range session.ReceivedChunks {
			payload, err := os.ReadFile(filepath.Join(dir, fmt.Sprintf("chunk-%08d.bin", index)))
			if err != nil {
				return nil, fmt.Errorf("read artifact transfer chunk: %w", err)
			}
			if len(payload) != record.SizeBytes || validator.BytesChecksum(payload) != record.Checksum {
				return nil, validationError("durable artifact chunk checksum mismatch")
			}
			record.Payload = payload
			session.ReceivedChunks[index] = record
		}
		out = append(out, domain.CopyArtifactTransferSession(session))
	}
	return out, nil
}

func (store *FileArtifactBodyStore) PutPayload(artifactID string, payload []byte) error {
	store.mu.Lock()
	defer store.mu.Unlock()
	return writeAtomicFile(store.payloadPath(artifactID), payload, 0o600)
}

func (store *FileArtifactBodyStore) GetPayload(artifactID string) ([]byte, error) {
	store.mu.Lock()
	defer store.mu.Unlock()
	payload, err := os.ReadFile(store.payloadPath(artifactID))
	if errors.Is(err, os.ErrNotExist) {
		return nil, repo.ErrNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("read artifact payload: %w", err)
	}
	return payload, nil
}

func (store *FileArtifactBodyStore) transferDir(transferID string) string {
	return filepath.Join(store.rootDir, "transfers", stableStorageKey(transferID))
}

func (store *FileArtifactBodyStore) payloadPath(artifactID string) string {
	return filepath.Join(store.rootDir, "payloads", stableStorageKey(artifactID)+".bin")
}

func stableStorageKey(value string) string {
	sum := sha256.Sum256([]byte(value))
	return hex.EncodeToString(sum[:])
}

func writeAtomicFile(path string, payload []byte, mode os.FileMode) error {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return fmt.Errorf("create durable body directory: %w", err)
	}
	tmp := path + ".tmp"
	file, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
	if err != nil {
		return fmt.Errorf("open durable body temporary file: %w", err)
	}
	if _, err := file.Write(payload); err != nil {
		_ = file.Close()
		_ = os.Remove(tmp)
		return fmt.Errorf("write durable body: %w", err)
	}
	if err := file.Sync(); err != nil {
		_ = file.Close()
		_ = os.Remove(tmp)
		return fmt.Errorf("sync durable body: %w", err)
	}
	if err := file.Close(); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("close durable body: %w", err)
	}
	if err := os.Rename(tmp, path); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("replace durable body: %w", err)
	}
	return nil
}