Commilitia-Drop/internal/relay/manager.go

package relay

import (
	"context"
	"errors"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"
)

const (
	DefaultMaxSessions        = 8
	DefaultMaxSessionsPerUser = 3
	DefaultSessionCap         = 64 << 20 // 64 MiB
	DefaultChunkLimit         = 1 << 20  // 1 MiB

	// DefaultIdleTTL is how long a registered relay session may sit without any
	// chunk/stream activity before the reaper tears it down. A fallback that
	// registered but never got used (both peers vanished) is reclaimed rather
	// than pinning a slot — and, once bytes flowed, its 64 MiB ring — forever.
	DefaultIdleTTL = 2 * time.Minute

	// DefaultReapPeriod is how often the reaper goroutine sweeps for idle sessions.
	DefaultReapPeriod = 30 * time.Second
)

var (
	ErrTooManySessions = errors.New("relay: too many active sessions")
	ErrForbidden       = errors.New("relay: caller is not a participant of this session")
	ErrNotRegistered   = errors.New("relay: session not registered")
)

// Session pairs a Ring with the metadata that proves who is allowed to use it.
// The ring is allocated lazily on the first Acquire so a registered-but-unused
// fallback costs only bookkeeping until bytes actually flow.
type Session struct {
	ID       string
	UserID   string
	Sender   string // device name of the legitimate sender (POSTs chunks)
	Receiver string // device name of the legitimate receiver (GETs the stream)

	capBytes int
	ring     *Ring // nil until first Acquire; guarded by Manager.mu

	bytesWritten atomic.Int64
	lastActivity atomic.Int64 // unix seconds; the reaper's idle clock
	readerActive atomic.Bool  // single-consumer guard (a second reader tears the ring)
}

func (s *Session) touch() { s.lastActivity.Store(time.Now().Unix()) }

func (s *Session) Write(ctx context.Context, p []byte) (int, error) {
	s.touch()
	n, err := s.ring.Write(ctx, p)
	if n > 0 {
		s.bytesWritten.Add(int64(n))
	}
	return n, err
}

func (s *Session) Read(ctx context.Context, p []byte) (int, error) {
	s.touch()
	return s.ring.Read(ctx, p)
}

func (s *Session) CloseWriter() {
	if s.ring != nil {
		s.ring.CloseWriter()
	}
}

func (s *Session) Abort() {
	if s.ring != nil {
		s.ring.Abort()
	}
}

func (s *Session) BytesWritten() int64 { return s.bytesWritten.Load() }

// AcquireReader claims the single consumer slot. The relay stream is strictly
// single-consumer: two concurrent readers would each drain bytes out of the
// ring and silently corrupt the receiver's copy (G2). A second reader gets a
// false here so the handler can 409 instead of joining.
func (s *Session) AcquireReader() bool { return s.readerActive.CompareAndSwap(false, true) }

// ReleaseReader frees the consumer slot when the stream handler returns.
func (s *Session) ReleaseReader() { s.readerActive.Store(false) }

// Manager owns active relay sessions. A session must be Register-ed (which the
// transfer service does when a transfer enters RELAY_ACTIVE) before either peer
// can Acquire it. This is the core of the R1/G2 hardening: the old Acquire
// minted a session for any unknown id on demand, so 8 one-byte requests could
// exhaust every slot and pin ~512 MiB. Now only legitimate, server-minted
// transfers that actually reached RELAY_ACTIVE get a slot, capped both globally
// and per-user, and reclaimed when idle.
type Manager struct {
	mu          sync.Mutex
	sessions    map[string]*Session
	perUser     map[string]int
	maxSessions int
	maxPerUser  int
	sessionCap  int
	idleTTL     time.Duration
}

func NewManager(maxSessions, sessionCapBytes int) *Manager {
	if maxSessions <= 0 {
		maxSessions = DefaultMaxSessions
	}
	if sessionCapBytes <= 0 {
		sessionCapBytes = DefaultSessionCap
	}
	maxPerUser := DefaultMaxSessionsPerUser
	if maxPerUser > maxSessions {
		// A per-user cap above the global cap is meaningless; clamp so small
		// global caps (tests, constrained deploys) still behave sanely.
		maxPerUser = maxSessions
	}
	return &Manager{
		sessions:    map[string]*Session{},
		perUser:     map[string]int{},
		maxSessions: maxSessions,
		maxPerUser:  maxPerUser,
		sessionCap:  sessionCapBytes,
		idleTTL:     DefaultIdleTTL,
	}
}

// Register reserves a relay slot for a transfer that has just entered
// RELAY_ACTIVE, recording its two legitimate participants. Only a registered
// session can later be joined via Acquire.
//
// Idempotent for the same (id, userID): a duplicate Register just refreshes the
// idle clock. A mismatched userID is ErrForbidden (ids are server-minted, so
// this is defensive). Returns ErrTooManySessions when the global or per-user
// cap is reached — the caller (transfer service) surfaces that as a retryable
// failure rather than moving the transfer into a relay mode with no slot.
func (m *Manager) Register(id, userID, sender, receiver string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.sessions[id]; ok {
		if s.UserID != userID {
			return ErrForbidden
		}
		s.touch()
		return nil
	}
	if len(m.sessions) >= m.maxSessions {
		return ErrTooManySessions
	}
	if m.perUser[userID] >= m.maxPerUser {
		return ErrTooManySessions
	}

	s := &Session{
		ID:       id,
		UserID:   userID,
		Sender:   sender,
		Receiver: receiver,
		capBytes: m.sessionCap,
	}
	s.touch()
	m.sessions[id] = s
	m.perUser[userID] = m.perUser[userID] + 1
	return nil
}

// Acquire joins an already-registered relay session as one of its two
// participants. Unlike the old implementation it never creates a session:
//   - unknown id          → ErrNotRegistered
//   - different user      → ErrForbidden
//   - not sender/receiver → ErrForbidden
//
// device is the caller's X-Device-Name (set by the auth middleware). The 64 MiB
// ring is allocated here, lazily, on the first join.
func (m *Manager) Acquire(id, userID, device string) (*Session, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	s, ok := m.sessions[id]
	if !ok {
		return nil, ErrNotRegistered
	}
	if s.UserID != userID {
		return nil, ErrForbidden
	}
	if device == "" || (device != s.Sender && device != s.Receiver) {
		return nil, ErrForbidden
	}
	if s.ring == nil {
		s.ring = newRing(s.capBytes)
	}
	s.touch()
	return s, nil
}

// Release removes the session from the active set and aborts any blocked
// Read/Write. Safe to call multiple times (terminal transitions and the pending
// sweeper both call it). A no-op for an id that was never registered.
func (m *Manager) Release(id string) {
	m.mu.Lock()
	s, ok := m.sessions[id]
	if ok {
		delete(m.sessions, id)
		m.decUserLocked(s.UserID)
	}
	m.mu.Unlock()
	if ok {
		s.Abort()
	}
}

// decUserLocked drops one from a user's active count, deleting the key at zero
// so perUser never accumulates empty entries. Caller holds m.mu.
func (m *Manager) decUserLocked(userID string) {
	if n := m.perUser[userID]; n <= 1 {
		delete(m.perUser, userID)
	} else {
		m.perUser[userID] = n - 1
	}
}

// Active reports the current count of held sessions.
func (m *Manager) Active() int {
	m.mu.Lock()
	defer m.mu.Unlock()
	return len(m.sessions)
}

// reapIdle aborts and removes every session whose last activity is older than
// the idle TTL, returning how many it reclaimed. now is unix seconds (injected
// so tests don't depend on the wall clock).
func (m *Manager) reapIdle(now int64) int {
	cutoff := now - int64(m.idleTTL/time.Second)
	var dead []*Session
	m.mu.Lock()
	for id, s := range m.sessions {
		if s.lastActivity.Load() < cutoff {
			delete(m.sessions, id)
			m.decUserLocked(s.UserID)
			dead = append(dead, s)
		}
	}
	m.mu.Unlock()
	for _, s := range dead {
		s.Abort()
	}
	return len(dead)
}

// RunReaper blocks until ctx is cancelled, periodically reclaiming idle relay
// sessions (fallbacks that registered but never completed — e.g. both peers
// went away mid-transfer). Mirrors transfer.RunSweeper; start it once at boot.
func (m *Manager) RunReaper(ctx context.Context) {
	t := time.NewTicker(DefaultReapPeriod)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			if n := m.reapIdle(time.Now().Unix()); n > 0 {
				slog.Info("relay reaper: reclaimed idle sessions", "count", n)
			}
		}
	}
}