silo-server/internal/uploads/session.go

package uploads

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"math"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const (
	defaultSessionTTL = 2 * time.Hour
	defaultMaxSize    = 256 << 20
	defaultChunkSize  = 8 << 20
)

var (
	ErrAlreadyCompleted = errors.New("upload session is already completed")
	ErrExpired          = errors.New("upload session expired")
	ErrIncomplete       = errors.New("upload session is incomplete")
	ErrInvalidChunk     = errors.New("invalid upload chunk")
	ErrInvalidRequest   = errors.New("invalid upload session request")
	ErrNotFound         = errors.New("upload session not found")
	ErrTooLarge         = errors.New("upload exceeds maximum size")
)

type ManagerOptions struct {
	RootDir      string
	TTL          time.Duration
	MaxSize      int64
	MaxChunkSize int64
	Now          func() time.Time
}

type CreateRequest struct {
	Filename  string
	SizeBytes int64
	ChunkSize int64
}

type SessionInfo struct {
	ID             string
	Filename       string
	SizeBytes      int64
	ChunkSize      int64
	TotalChunks    int
	ReceivedChunks int
	ReceivedBytes  int64
	Complete       bool
	ExpiresAt      time.Time
}

type CompletedUpload struct {
	ID        string
	Filename  string
	SizeBytes int64
	Path      string
	Cleanup   func()
}

type Manager struct {
	mu           sync.Mutex
	rootDir      string
	ttl          time.Duration
	maxSize      int64
	maxChunkSize int64
	now          func() time.Time
	sessions     map[string]*session
}

type session struct {
	id             string
	filename       string
	sizeBytes      int64
	chunkSize      int64
	totalChunks    int
	received       []bool
	receivedChunks int
	receivedBytes  int64
	path           string
	dir            string
	expiresAt      time.Time
	completed      bool
}

func NewManager(opts ManagerOptions) *Manager {
	rootDir := opts.RootDir
	if strings.TrimSpace(rootDir) == "" {
		rootDir = filepath.Join(os.TempDir(), "silo-uploads")
	}
	ttl := opts.TTL
	if ttl <= 0 {
		ttl = defaultSessionTTL
	}
	maxSize := opts.MaxSize
	if maxSize <= 0 {
		maxSize = defaultMaxSize
	}
	maxChunkSize := opts.MaxChunkSize
	if maxChunkSize <= 0 {
		maxChunkSize = defaultChunkSize
	}
	now := opts.Now
	if now == nil {
		now = time.Now
	}
	return &Manager{
		rootDir:      rootDir,
		ttl:          ttl,
		maxSize:      maxSize,
		maxChunkSize: maxChunkSize,
		now:          now,
		sessions:     make(map[string]*session),
	}
}

func (m *Manager) MaxChunkSize() int64 {
	return m.maxChunkSize
}

func (m *Manager) Create(req CreateRequest) (SessionInfo, error) {
	filename := filepath.Base(strings.TrimSpace(req.Filename))
	if filename == "." || filename == string(filepath.Separator) || filename == "" {
		filename = "upload.bin"
	}
	if req.SizeBytes <= 0 {
		return SessionInfo{}, fmt.Errorf("%w: size_bytes must be positive", ErrInvalidRequest)
	}
	if req.SizeBytes > m.maxSize {
		return SessionInfo{}, fmt.Errorf("%w: maximum size is %d bytes", ErrTooLarge, m.maxSize)
	}

	chunkSize := req.ChunkSize
	if chunkSize <= 0 {
		chunkSize = m.maxChunkSize
	}
	if chunkSize > m.maxChunkSize {
		return SessionInfo{}, fmt.Errorf("%w: chunk_size must not exceed %d bytes", ErrInvalidRequest, m.maxChunkSize)
	}

	totalChunks64 := (req.SizeBytes + chunkSize - 1) / chunkSize
	if totalChunks64 <= 0 || totalChunks64 > math.MaxInt32 {
		return SessionInfo{}, fmt.Errorf("%w: invalid total chunk count", ErrInvalidRequest)
	}

	if err := os.MkdirAll(m.rootDir, 0o700); err != nil {
		return SessionInfo{}, fmt.Errorf("create upload root: %w", err)
	}

	id, err := newID()
	if err != nil {
		return SessionInfo{}, err
	}
	dir := filepath.Join(m.rootDir, id)
	if err := os.Mkdir(dir, 0o700); err != nil {
		return SessionInfo{}, fmt.Errorf("create upload session directory: %w", err)
	}

	path := filepath.Join(dir, "upload.bin")
	file, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_RDWR, 0o600)
	if err != nil {
		_ = os.RemoveAll(dir)
		return SessionInfo{}, fmt.Errorf("create upload session file: %w", err)
	}
	if err := file.Truncate(req.SizeBytes); err != nil {
		_ = file.Close()
		_ = os.RemoveAll(dir)
		return SessionInfo{}, fmt.Errorf("size upload session file: %w", err)
	}
	if err := file.Close(); err != nil {
		_ = os.RemoveAll(dir)
		return SessionInfo{}, fmt.Errorf("close upload session file: %w", err)
	}

	now := m.now()
	s := &session{
		id:          id,
		filename:    filename,
		sizeBytes:   req.SizeBytes,
		chunkSize:   chunkSize,
		totalChunks: int(totalChunks64),
		received:    make([]bool, int(totalChunks64)),
		path:        path,
		dir:         dir,
		expiresAt:   now.Add(m.ttl),
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	m.cleanupExpiredLocked(now)
	m.sessions[id] = s
	return s.info(), nil
}

func (m *Manager) PutChunk(ctx context.Context, id string, index int, body io.Reader, contentLength int64) (SessionInfo, error) {
	if body == nil {
		return SessionInfo{}, fmt.Errorf("%w: chunk body is required", ErrInvalidChunk)
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	m.cleanupExpiredLocked(m.now())

	s, err := m.getLocked(id)
	if err != nil {
		return SessionInfo{}, err
	}
	if s.completed {
		return SessionInfo{}, ErrAlreadyCompleted
	}
	if index < 0 || index >= s.totalChunks {
		return SessionInfo{}, fmt.Errorf("%w: chunk index is out of range", ErrInvalidChunk)
	}

	expectedSize := s.expectedChunkSize(index)
	if contentLength >= 0 && contentLength != expectedSize {
		return SessionInfo{}, fmt.Errorf("%w: chunk size must be %d bytes", ErrInvalidChunk, expectedSize)
	}
	if expectedSize > m.maxChunkSize {
		return SessionInfo{}, fmt.Errorf("%w: maximum chunk size is %d bytes", ErrTooLarge, m.maxChunkSize)
	}
	if s.received[index] {
		return s.info(), nil
	}

	file, err := os.OpenFile(s.path, os.O_WRONLY, 0o600)
	if err != nil {
		return SessionInfo{}, fmt.Errorf("open upload session file: %w", err)
	}
	defer file.Close()

	writer := io.NewOffsetWriter(file, int64(index)*s.chunkSize)
	written, err := io.Copy(writer, io.LimitReader(body, expectedSize))
	if err != nil {
		return SessionInfo{}, err
	}
	if err := ctx.Err(); err != nil {
		return SessionInfo{}, err
	}
	if written != expectedSize {
		return SessionInfo{}, fmt.Errorf("%w: chunk size must be %d bytes", ErrInvalidChunk, expectedSize)
	}
	var extra [1]byte
	extraBytes, err := body.Read(extra[:])
	if err != nil && !errors.Is(err, io.EOF) {
		return SessionInfo{}, err
	}
	if extraBytes > 0 {
		return SessionInfo{}, fmt.Errorf("%w: chunk size must be %d bytes", ErrInvalidChunk, expectedSize)
	}

	s.received[index] = true
	s.receivedChunks++
	s.receivedBytes += written
	return s.info(), nil
}

func (m *Manager) Complete(id string) (*CompletedUpload, error) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.cleanupExpiredLocked(m.now())

	s, err := m.getLocked(id)
	if err != nil {
		return nil, err
	}
	if s.completed {
		return nil, ErrAlreadyCompleted
	}
	if !s.isComplete() {
		return nil, ErrIncomplete
	}
	s.completed = true
	delete(m.sessions, id)

	return &CompletedUpload{
		ID:        s.id,
		Filename:  s.filename,
		SizeBytes: s.sizeBytes,
		Path:      s.path,
		Cleanup: func() {
			_ = os.RemoveAll(s.dir)
		},
	}, nil
}

func (m *Manager) Cancel(id string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	s, ok := m.sessions[id]
	if !ok {
		return ErrNotFound
	}
	delete(m.sessions, id)
	return os.RemoveAll(s.dir)
}

func (m *Manager) getLocked(id string) (*session, error) {
	id = strings.TrimSpace(id)
	if id == "" {
		return nil, ErrNotFound
	}
	s, ok := m.sessions[id]
	if !ok {
		return nil, ErrNotFound
	}
	if !m.now().Before(s.expiresAt) {
		delete(m.sessions, id)
		_ = os.RemoveAll(s.dir)
		return nil, ErrExpired
	}
	return s, nil
}

func (m *Manager) cleanupExpiredLocked(now time.Time) {
	for id, s := range m.sessions {
		if now.Before(s.expiresAt) {
			continue
		}
		delete(m.sessions, id)
		_ = os.RemoveAll(s.dir)
	}
}

func (s *session) isComplete() bool {
	return s.receivedChunks == s.totalChunks && s.receivedBytes == s.sizeBytes
}

func (s *session) expectedChunkSize(index int) int64 {
	offset := int64(index) * s.chunkSize
	remaining := s.sizeBytes - offset
	if remaining < s.chunkSize {
		return remaining
	}
	return s.chunkSize
}

func (s *session) info() SessionInfo {
	return SessionInfo{
		ID:             s.id,
		Filename:       s.filename,
		SizeBytes:      s.sizeBytes,
		ChunkSize:      s.chunkSize,
		TotalChunks:    s.totalChunks,
		ReceivedChunks: s.receivedChunks,
		ReceivedBytes:  s.receivedBytes,
		Complete:       s.receivedChunks == s.totalChunks && s.receivedBytes == s.sizeBytes,
		ExpiresAt:      s.expiresAt,
	}
}

func newID() (string, error) {
	var b [16]byte
	if _, err := rand.Read(b[:]); err != nil {
		return "", fmt.Errorf("generate upload session id: %w", err)
	}
	return hex.EncodeToString(b[:]), nil
}