browser/platform/service/artifact_body_store.go

package service

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"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)
	ReadPayloadRange(string, int64, int) ([]byte, error)
	CommitTransferPayload(domain.ArtifactTransferSession) 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
}

func (store *MemoryArtifactBodyStore) ReadPayloadRange(artifactID string, offset int64, length int) ([]byte, error) {
	store.mu.Lock()
	defer store.mu.Unlock()
	payload, exists := store.payloads[artifactID]
	if !exists {
		return nil, repo.ErrNotFound
	}
	if offset < 0 || offset > int64(len(payload)) || length < 0 || int64(length) > int64(len(payload))-offset {
		return nil, fmt.Errorf("artifact range is invalid")
	}
	return domain.CopyBytes(payload[int(offset) : int(offset)+length]), nil
}

func (store *MemoryArtifactBodyStore) CommitTransferPayload(session domain.ArtifactTransferSession) error {
	store.mu.Lock()
	defer store.mu.Unlock()
	payload := make([]byte, 0, int(session.SizeBytes))
	for index := 0; index < session.TotalChunks; index++ {
		record, exists := session.ReceivedChunks[index]
		if !exists {
			return validationError("artifact transfer has missing chunks")
		}
		payload = append(payload, record.Payload...)
	}
	if int64(len(payload)) != session.SizeBytes {
		return validationError("artifact transfer size does not match metadata")
	}
	if checksum := validator.BytesChecksum(payload); checksum != session.Checksum {
		return validationError("artifact transfer checksum does not match metadata")
	}
	store.payloads[session.ArtifactID] = payload
	return 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 {
		if record.Payload != nil {
			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 {
			chunkPath := filepath.Join(dir, fmt.Sprintf("chunk-%08d.bin", index))
			if _, err := os.Stat(chunkPath); err != nil {
				return nil, fmt.Errorf("stat artifact transfer chunk: %w", err)
			}
			record.Payload = nil
			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) ReadPayloadRange(artifactID string, offset int64, length int) ([]byte, error) {
	if offset < 0 || length < 0 {
		return nil, fmt.Errorf("artifact range is invalid")
	}
	store.mu.Lock()
	defer store.mu.Unlock()
	file, err := os.Open(store.payloadPath(artifactID))
	if errors.Is(err, os.ErrNotExist) {
		return nil, repo.ErrNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("open artifact payload: %w", err)
	}
	defer file.Close()
	payload := make([]byte, length)
	read, err := file.ReadAt(payload, offset)
	if err != nil && !(errors.Is(err, io.ErrUnexpectedEOF) && read == length) {
		return nil, fmt.Errorf("read artifact payload range: %w", err)
	}
	if read != length {
		return nil, fmt.Errorf("artifact payload range is shorter than requested")
	}
	return payload, nil
}

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

	transferDir := store.transferDir(session.TransferID)
	payloadPath := store.payloadPath(session.ArtifactID)
	if err := os.MkdirAll(filepath.Dir(payloadPath), 0o700); err != nil {
		return fmt.Errorf("create artifact payload directory: %w", err)
	}
	tmp := payloadPath + ".tmp"
	out, err := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open artifact payload temporary file: %w", err)
	}
	hash := sha256.New()
	written := int64(0)
	for index := 0; index < session.TotalChunks; index++ {
		record, exists := session.ReceivedChunks[index]
		if !exists {
			_ = out.Close()
			_ = os.Remove(tmp)
			return validationError("artifact transfer has missing chunks")
		}
		chunkPath := filepath.Join(transferDir, fmt.Sprintf("chunk-%08d.bin", index))
		chunk, err := os.Open(chunkPath)
		if err != nil {
			_ = out.Close()
			_ = os.Remove(tmp)
			return fmt.Errorf("open artifact transfer chunk: %w", err)
		}
		chunkHash := sha256.New()
		count, copyErr := io.Copy(io.MultiWriter(out, hash, chunkHash), chunk)
		closeErr := chunk.Close()
		if copyErr != nil {
			_ = out.Close()
			_ = os.Remove(tmp)
			return fmt.Errorf("copy artifact transfer chunk: %w", copyErr)
		}
		if closeErr != nil {
			_ = out.Close()
			_ = os.Remove(tmp)
			return fmt.Errorf("close artifact transfer chunk: %w", closeErr)
		}
		if int(count) != record.SizeBytes || "sha256:"+hex.EncodeToString(chunkHash.Sum(nil)) != record.Checksum {
			_ = out.Close()
			_ = os.Remove(tmp)
			return validationError("durable artifact chunk checksum mismatch")
		}
		written += count
	}
	if written != session.SizeBytes {
		_ = out.Close()
		_ = os.Remove(tmp)
		return validationError("artifact transfer size does not match metadata")
	}
	if checksum := "sha256:" + hex.EncodeToString(hash.Sum(nil)); checksum != session.Checksum {
		_ = out.Close()
		_ = os.Remove(tmp)
		return validationError("artifact transfer checksum does not match metadata")
	}
	if err := out.Sync(); err != nil {
		_ = out.Close()
		_ = os.Remove(tmp)
		return fmt.Errorf("sync artifact payload: %w", err)
	}
	if err := out.Close(); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("close artifact payload: %w", err)
	}
	if err := os.Rename(tmp, payloadPath); err != nil {
		_ = os.Remove(tmp)
		return fmt.Errorf("commit artifact payload: %w", err)
	}
	return 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
}