soft_quay/core/storage/download_task.go

package storage

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"

	"softbox.local/core/downloader"
)

const (
	maxDownloadTaskDocumentBytes = 64 * 1024
	maxDownloadTaskCount         = 4096
)

var ErrDownloadTaskLayoutUnsafe = errors.New("download task storage layout is unsafe")

// DownloadTaskStore atomically persists download metadata below
// downloads/tasks. Transfer bytes live separately below downloads/files.
type DownloadTaskStore struct {
	root string
	mu   sync.Mutex
}

// NewDownloadTaskStore creates a strict JSON task store.
func NewDownloadTaskStore(downloadsRoot string) *DownloadTaskStore {
	return &DownloadTaskStore{root: downloadsRoot}
}

// Write validates and atomically replaces one task document.
func (store *DownloadTaskStore) Write(task downloader.Task) error {
	store.mu.Lock()
	defer store.mu.Unlock()

	if err := task.Validate(); err != nil {
		return err
	}
	paths, err := store.ensureLayout(task.RequestID)
	if err != nil {
		return err
	}
	document, err := json.Marshal(task)
	if err != nil {
		return fmt.Errorf("encode download task: %w", err)
	}
	document = append(document, '\n')
	if len(document) > maxDownloadTaskDocumentBytes {
		return fmt.Errorf("%w: metadata is too large", downloader.ErrInvalidTask)
	}
	return replaceDownloadTaskFile(paths.TasksDir, paths.Metadata, paths.Backup, document)
}

// LoadAll reads every task, falling back to a backup only when its main file
// is absent.
func (store *DownloadTaskStore) LoadAll() ([]downloader.Task, error) {
	store.mu.Lock()
	defer store.mu.Unlock()

	paths, err := store.ensureLayout("layout-probe")
	if err != nil {
		return nil, err
	}
	entries, err := os.ReadDir(paths.TasksDir)
	if err != nil {
		return nil, fmt.Errorf("read download task directory: %w", err)
	}
	requestIDs := make(map[string]struct{})
	for _, entry := range entries {
		name := entry.Name()
		if strings.HasPrefix(name, ".download-task-") && strings.HasSuffix(name, ".tmp") {
			continue
		}
		requestID, recognized := metadataRequestID(name)
		if !recognized || !downloader.ValidRequestID(requestID) {
			return nil, fmt.Errorf(
				"%w: unexpected metadata entry %q",
				ErrDownloadTaskLayoutUnsafe,
				name,
			)
		}
		info, infoErr := entry.Info()
		if infoErr != nil {
			return nil, fmt.Errorf("inspect download task entry: %w", infoErr)
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
			return nil, fmt.Errorf(
				"%w: metadata entry is not regular",
				ErrDownloadTaskLayoutUnsafe,
			)
		}
		requestIDs[requestID] = struct{}{}
	}
	if len(requestIDs) > maxDownloadTaskCount {
		return nil, fmt.Errorf("%w: too many tasks", downloader.ErrInvalidTask)
	}
	sortedIDs := make([]string, 0, len(requestIDs))
	for requestID := range requestIDs {
		sortedIDs = append(sortedIDs, requestID)
	}
	sort.Strings(sortedIDs)

	tasks := make([]downloader.Task, 0, len(sortedIDs))
	for _, requestID := range sortedIDs {
		taskPaths, deriveErr := downloader.DeriveTaskPaths(store.root, requestID)
		if deriveErr != nil {
			return nil, deriveErr
		}
		document, readErr := readBoundedRegularFile(
			taskPaths.Metadata,
			maxDownloadTaskDocumentBytes,
		)
		if os.IsNotExist(readErr) {
			document, readErr = readBoundedRegularFile(
				taskPaths.Backup,
				maxDownloadTaskDocumentBytes,
			)
		}
		if readErr != nil {
			return nil, readErr
		}
		task, decodeErr := decodeDownloadTask(document)
		if decodeErr != nil {
			return nil, decodeErr
		}
		if task.RequestID != requestID {
			return nil, fmt.Errorf(
				"%w: request_id does not match metadata filename",
				downloader.ErrInvalidTask,
			)
		}
		tasks = append(tasks, task)
	}
	return tasks, nil
}

// Delete removes the exact metadata and backup files for a request.
func (store *DownloadTaskStore) Delete(requestID string) error {
	store.mu.Lock()
	defer store.mu.Unlock()

	paths, err := downloader.DeriveTaskPaths(store.root, requestID)
	if err != nil {
		return err
	}
	for _, candidate := range []string{paths.Metadata, paths.Backup} {
		info, inspectErr := os.Lstat(candidate)
		if os.IsNotExist(inspectErr) {
			continue
		}
		if inspectErr != nil {
			return fmt.Errorf("inspect download task metadata: %w", inspectErr)
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
			return fmt.Errorf(
				"%w: metadata is not a regular file",
				ErrDownloadTaskLayoutUnsafe,
			)
		}
		if removeErr := os.Remove(candidate); removeErr != nil {
			return fmt.Errorf("remove download task metadata: %w", removeErr)
		}
	}
	return nil
}

func (store *DownloadTaskStore) ensureLayout(requestID string) (downloader.TaskPaths, error) {
	paths, err := downloader.DeriveTaskPaths(store.root, requestID)
	if err != nil {
		return downloader.TaskPaths{}, err
	}
	if err := os.MkdirAll(paths.Root, 0o700); err != nil {
		return downloader.TaskPaths{}, fmt.Errorf("create downloads root: %w", err)
	}
	for _, directory := range []string{paths.Root, paths.TasksDir, paths.FilesDir} {
		if err := os.Mkdir(directory, 0o700); err != nil && !os.IsExist(err) {
			return downloader.TaskPaths{}, fmt.Errorf("create download directory: %w", err)
		}
		info, inspectErr := os.Lstat(directory)
		if inspectErr != nil {
			return downloader.TaskPaths{}, fmt.Errorf("inspect download directory: %w", inspectErr)
		}
		if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
			return downloader.TaskPaths{}, fmt.Errorf(
				"%w: managed path is not a real directory",
				ErrDownloadTaskLayoutUnsafe,
			)
		}
	}
	return paths, nil
}

func metadataRequestID(name string) (string, bool) {
	switch {
	case strings.HasSuffix(name, ".json.backup"):
		return strings.TrimSuffix(name, ".json.backup"), true
	case strings.HasSuffix(name, ".json"):
		return strings.TrimSuffix(name, ".json"), true
	default:
		return "", false
	}
}

func decodeDownloadTask(document []byte) (downloader.Task, error) {
	var task downloader.Task
	decoder := json.NewDecoder(bytes.NewReader(document))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&task); err != nil {
		return downloader.Task{}, fmt.Errorf("%w: decode: %v", downloader.ErrInvalidTask, err)
	}
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		if err == nil {
			return downloader.Task{}, fmt.Errorf(
				"%w: trailing JSON value",
				downloader.ErrInvalidTask,
			)
		}
		return downloader.Task{}, fmt.Errorf(
			"%w: trailing data: %v",
			downloader.ErrInvalidTask,
			err,
		)
	}
	if err := task.Validate(); err != nil {
		return downloader.Task{}, err
	}
	return task, nil
}

func readBoundedRegularFile(filePath string, maximum int64) ([]byte, error) {
	info, err := os.Lstat(filePath)
	if err != nil {
		return nil, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return nil, fmt.Errorf(
			"%w: metadata is not a regular file",
			ErrDownloadTaskLayoutUnsafe,
		)
	}
	if info.Size() > maximum {
		return nil, fmt.Errorf("%w: metadata is too large", downloader.ErrInvalidTask)
	}
	file, err := os.Open(filePath)
	if err != nil {
		return nil, fmt.Errorf("open download task metadata: %w", err)
	}
	defer file.Close()
	document, err := io.ReadAll(io.LimitReader(file, maximum+1))
	if err != nil {
		return nil, fmt.Errorf("read download task metadata: %w", err)
	}
	if int64(len(document)) > maximum {
		return nil, fmt.Errorf("%w: metadata is too large", downloader.ErrInvalidTask)
	}
	return document, nil
}

func replaceDownloadTaskFile(directory, target, backup string, document []byte) error {
	temporary, err := os.CreateTemp(directory, ".download-task-*.tmp")
	if err != nil {
		return fmt.Errorf("create download task temp file: %w", err)
	}
	temporaryPath := temporary.Name()
	defer os.Remove(temporaryPath)

	if err := temporary.Chmod(0o600); err != nil {
		temporary.Close()
		return fmt.Errorf("protect download task temp file: %w", err)
	}
	if _, err := temporary.Write(document); err != nil {
		temporary.Close()
		return fmt.Errorf("write download task temp file: %w", err)
	}
	if err := temporary.Sync(); err != nil {
		temporary.Close()
		return fmt.Errorf("sync download task temp file: %w", err)
	}
	if err := temporary.Close(); err != nil {
		return fmt.Errorf("close download task temp file: %w", err)
	}

	movedTarget := false
	hadBackup := false
	if info, inspectErr := os.Lstat(target); inspectErr == nil {
		if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
			return fmt.Errorf(
				"%w: metadata target is not regular",
				ErrDownloadTaskLayoutUnsafe,
			)
		}
		if err := removeExactRegularFile(backup); err != nil {
			return err
		}
		if err := os.Rename(target, backup); err != nil {
			return fmt.Errorf("backup download task metadata: %w", err)
		}
		movedTarget = true
	} else if !os.IsNotExist(inspectErr) {
		return fmt.Errorf("inspect download task metadata: %w", inspectErr)
	} else {
		backupInfo, backupErr := os.Lstat(backup)
		switch {
		case backupErr == nil:
			if backupInfo.Mode()&os.ModeSymlink != 0 || !backupInfo.Mode().IsRegular() {
				return fmt.Errorf(
					"%w: metadata backup is not regular",
					ErrDownloadTaskLayoutUnsafe,
				)
			}
			hadBackup = true
		case !os.IsNotExist(backupErr):
			return fmt.Errorf("inspect download task backup: %w", backupErr)
		}
	}
	if err := os.Rename(temporaryPath, target); err != nil {
		if movedTarget {
			_ = os.Rename(backup, target)
		}
		return fmt.Errorf("activate download task metadata: %w", err)
	}
	if movedTarget || hadBackup {
		if err := removeExactRegularFile(backup); err != nil {
			return err
		}
	}
	return nil
}

func removeExactRegularFile(filePath string) error {
	info, err := os.Lstat(filePath)
	if os.IsNotExist(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("inspect download task backup: %w", err)
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return fmt.Errorf(
			"%w: metadata backup is not regular",
			ErrDownloadTaskLayoutUnsafe,
		)
	}
	if err := os.Remove(filePath); err != nil {
		return fmt.Errorf("remove download task backup: %w", err)
	}
	return nil
}