teraWeb/backend/internal/directstream/ticket.go

package directstream

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"sync"
	"time"
)

const (
	TicketEntropyBytes = 32
	defaultTicketTTL   = 30 * time.Second
	defaultMaxPending  = 256
)

var (
	ErrInvalidTicketClaims = errors.New("invalid direct stream ticket claims")
	ErrInvalidTicket       = errors.New("invalid or expired direct stream ticket")
	ErrTicketCapacity      = errors.New("direct stream ticket capacity reached")
)

type TicketClaims struct {
	SessionID string
	Channel   string
	Actor     string
	Access    string
}

type TicketBinding struct {
	SessionID string
	Channel   string
}

type IssuedTicket struct {
	Value     string
	ExpiresAt time.Time
}

type TicketStoreOptions struct {
	Clock      func() time.Time
	Entropy    io.Reader
	TTL        time.Duration
	MaxPending int
}

type ticketRecord struct {
	claims    TicketClaims
	expiresAt time.Time
}

type TicketStore struct {
	mu         sync.Mutex
	clock      func() time.Time
	entropy    io.Reader
	ttl        time.Duration
	maxPending int
	records    map[[sha256.Size]byte]ticketRecord
}

func NewTicketStore(options TicketStoreOptions) *TicketStore {
	clock := options.Clock
	if clock == nil {
		clock = time.Now
	}
	entropy := options.Entropy
	if entropy == nil {
		entropy = rand.Reader
	}
	ttl := options.TTL
	if ttl <= 0 {
		ttl = defaultTicketTTL
	}
	maxPending := options.MaxPending
	if maxPending <= 0 {
		maxPending = defaultMaxPending
	}
	return &TicketStore{
		clock:      clock,
		entropy:    entropy,
		ttl:        ttl,
		maxPending: maxPending,
		records:    make(map[[sha256.Size]byte]ticketRecord),
	}
}

func (s *TicketStore) Issue(claims TicketClaims) (IssuedTicket, error) {
	if err := validateTicketClaims(claims); err != nil {
		return IssuedTicket{}, err
	}
	now := s.clock().UTC()
	s.mu.Lock()
	defer s.mu.Unlock()
	s.pruneExpiredLocked(now)
	if len(s.records) >= s.maxPending {
		return IssuedTicket{}, ErrTicketCapacity
	}

	raw := make([]byte, TicketEntropyBytes)
	if _, err := io.ReadFull(s.entropy, raw); err != nil {
		return IssuedTicket{}, fmt.Errorf("generate direct stream ticket: %w", err)
	}
	value := base64.RawURLEncoding.EncodeToString(raw)
	digest := sha256.Sum256([]byte(value))
	if _, exists := s.records[digest]; exists {
		return IssuedTicket{}, errors.New("generate direct stream ticket: entropy collision")
	}
	expiresAt := now.Add(s.ttl)
	s.records[digest] = ticketRecord{claims: claims, expiresAt: expiresAt}
	return IssuedTicket{Value: value, ExpiresAt: expiresAt}, nil
}

func (s *TicketStore) Consume(value string, binding TicketBinding) (TicketClaims, error) {
	if value == "" || binding.SessionID == "" || ValidateChannel(binding.Channel) != nil {
		return TicketClaims{}, ErrInvalidTicket
	}
	raw, err := base64.RawURLEncoding.DecodeString(value)
	if err != nil || len(raw) != TicketEntropyBytes {
		return TicketClaims{}, ErrInvalidTicket
	}
	digest := sha256.Sum256([]byte(value))
	now := s.clock().UTC()

	s.mu.Lock()
	defer s.mu.Unlock()
	record, ok := s.records[digest]
	if !ok || !now.Before(record.expiresAt) {
		delete(s.records, digest)
		return TicketClaims{}, ErrInvalidTicket
	}
	if record.claims.SessionID != binding.SessionID || record.claims.Channel != binding.Channel {
		return TicketClaims{}, ErrInvalidTicket
	}
	delete(s.records, digest)
	return record.claims, nil
}

func (s *TicketStore) pruneExpiredLocked(now time.Time) {
	for digest, record := range s.records {
		if !now.Before(record.expiresAt) {
			delete(s.records, digest)
		}
	}
}

func validateTicketClaims(claims TicketClaims) error {
	if claims.SessionID == "" || claims.Actor == "" || ValidateChannel(claims.Channel) != nil || ValidateAccess(claims.Access) != nil {
		return ErrInvalidTicketClaims
	}
	return nil
}