teraWeb/backend/internal/directstream/ticket_socket.go

package directstream

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"
)

const (
	DefaultTicketSocketPath = "/run/tera-api/direct-tickets.sock"
	maxTicketSocketBytes    = 8 * 1024
	defaultTicketSocketTTL  = 10 * time.Second
	defaultActiveBindings   = 256
	defaultActiveBindingTTL = 60 * time.Second
)

var ErrTicketBindingInactive = errors.New("direct stream ticket binding is not active")

type TicketConsumer interface {
	ConsumeTicket(context.Context, string, TicketBinding) (TicketClaims, error)
}

type TicketLifecycle interface {
	CloseDirectStream(context.Context, TicketBinding) error
	ReleaseDirectWriter(context.Context, TicketBinding) error
	HeartbeatDirectStream(context.Context, TicketBinding) error
	RecordDirectStreamEvent(context.Context, TicketBinding, SessionEvent) error
}

type TicketAuthority interface {
	TicketConsumer
	CloseConsumedStream(context.Context, TicketClaims) error
	ReleaseConsumedWriter(context.Context, TicketClaims) error
	RecordConsumedStreamEvent(context.Context, TicketClaims, SessionEvent) error
}

type ticketSocketRequest struct {
	Action    string        `json:"action"`
	Ticket    string        `json:"ticket,omitempty"`
	SessionID string        `json:"session_id"`
	Channel   string        `json:"channel"`
	Event     *SessionEvent `json:"event,omitempty"`
}

type ticketSocketResponse struct {
	Claims *TicketClaims          `json:"claims,omitempty"`
	Status *TicketServiceSnapshot `json:"status,omitempty"`
	Code   string                 `json:"code,omitempty"`
	Error  string                 `json:"error,omitempty"`
}

type TicketServiceSnapshot struct {
	ActiveBindings       int    `json:"active_bindings"`
	MaxActiveBindings    int    `json:"max_active_bindings"`
	Consumed             uint64 `json:"consumed"`
	Heartbeats           uint64 `json:"heartbeats"`
	ReleasedWriters      uint64 `json:"released_writers"`
	ClosedStreams        uint64 `json:"closed_streams"`
	ExpiredStreams       uint64 `json:"expired_streams"`
	LifecycleFailures    uint64 `json:"lifecycle_failures"`
	SessionEvents        uint64 `json:"session_events"`
	SessionEventFailures uint64 `json:"session_event_failures"`
}

type activeTicketBinding struct {
	claims   TicketClaims
	lastSeen time.Time
}

// TicketSocketServer is the process boundary between tera-api and
// tera-sessiond. Lifecycle requests can refer only to claims successfully
// consumed on this server; the client cannot forge actor or access fields.
type TicketSocketServer struct {
	Authority       TicketAuthority
	MaxActiveClaims int
	Clock           func() time.Time
	ActiveTTL       time.Duration

	mu                   sync.Mutex
	active               map[TicketBinding]activeTicketBinding
	consumed             uint64
	heartbeats           uint64
	releasedWriters      uint64
	closedStreams        uint64
	expiredStreams       uint64
	lifecycleFailures    uint64
	sessionEvents        uint64
	sessionEventFailures uint64
}

func (s *TicketSocketServer) Serve(listener net.Listener) error {
	if s.Authority == nil {
		return errors.New("direct stream ticket authority is required")
	}
	s.mu.Lock()
	if s.active == nil {
		s.active = make(map[TicketBinding]activeTicketBinding)
	}
	s.mu.Unlock()
	stopJanitor := make(chan struct{})
	defer close(stopJanitor)
	go s.runJanitor(stopJanitor)
	for {
		connection, err := listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return nil
			}
			if temporary, ok := err.(interface{ Temporary() bool }); ok && temporary.Temporary() {
				continue
			}
			return err
		}
		go s.handleConnection(connection)
	}
}

func (s *TicketSocketServer) handleConnection(connection net.Conn) {
	defer connection.Close()
	_ = connection.SetDeadline(time.Now().Add(defaultTicketSocketTTL))
	payload, err := bufio.NewReader(io.LimitReader(connection, maxTicketSocketBytes+1)).ReadBytes('\n')
	if err != nil || len(payload) > maxTicketSocketBytes || !bytes.HasSuffix(payload, []byte{'\n'}) {
		s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "invalid direct stream ticket request"})
		return
	}
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()
	var request ticketSocketRequest
	if err := decoder.Decode(&request); err != nil {
		s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "invalid direct stream ticket request"})
		return
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "direct stream ticket request must contain one JSON value"})
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	if request.Action == "status" {
		if request.Ticket != "" || request.SessionID != "" || request.Channel != "" || request.Event != nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "status request must not include a ticket binding"})
			return
		}
		s.writeResponse(connection, ticketSocketResponse{Status: pointerToTicketStatus(s.Snapshot())})
		return
	}
	binding := TicketBinding{SessionID: request.SessionID, Channel: request.Channel}
	if binding.SessionID == "" || ValidateChannel(binding.Channel) != nil {
		s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "invalid direct stream ticket binding"})
		return
	}
	switch request.Action {
	case "consume":
		if request.Ticket == "" || request.Event != nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "direct stream ticket is required"})
			return
		}
		claims, err := s.consume(ctx, request.Ticket, binding)
		if err != nil {
			s.writeError(connection, err)
			return
		}
		s.writeResponse(connection, ticketSocketResponse{Claims: &claims})
	case "release_writer":
		if request.Ticket != "" || request.Event != nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "ticket is not accepted for lifecycle requests"})
			return
		}
		if err := s.releaseWriter(ctx, binding); err != nil {
			s.writeError(connection, err)
			return
		}
		s.writeResponse(connection, ticketSocketResponse{})
	case "heartbeat":
		if request.Ticket != "" || request.Event != nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "ticket is not accepted for lifecycle requests"})
			return
		}
		if err := s.heartbeat(binding); err != nil {
			s.writeError(connection, err)
			return
		}
		s.writeResponse(connection, ticketSocketResponse{})
	case "event":
		if request.Ticket != "" || request.Event == nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "session event request is invalid"})
			return
		}
		if err := s.recordEvent(ctx, binding, *request.Event); err != nil {
			s.writeError(connection, err)
			return
		}
		s.writeResponse(connection, ticketSocketResponse{})
	case "close":
		if request.Ticket != "" || request.Event != nil {
			s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "ticket is not accepted for lifecycle requests"})
			return
		}
		if err := s.close(ctx, binding); err != nil {
			s.writeError(connection, err)
			return
		}
		s.writeResponse(connection, ticketSocketResponse{})
	default:
		s.writeResponse(connection, ticketSocketResponse{Code: "validation_failed", Error: "direct stream ticket action is not supported"})
	}
}

func (s *TicketSocketServer) consume(ctx context.Context, ticket string, binding TicketBinding) (TicketClaims, error) {
	s.ExpireInactive(ctx)
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.active[binding]; exists {
		return TicketClaims{}, ErrInvalidTicket
	}
	maxActive := s.MaxActiveClaims
	if maxActive <= 0 {
		maxActive = defaultActiveBindings
	}
	if len(s.active) >= maxActive {
		return TicketClaims{}, ErrTicketCapacity
	}
	claims, err := s.Authority.ConsumeTicket(ctx, ticket, binding)
	if err != nil {
		return TicketClaims{}, err
	}
	s.active[binding] = activeTicketBinding{claims: claims, lastSeen: s.now()}
	s.consumed++
	return claims, nil
}

func (s *TicketSocketServer) releaseWriter(ctx context.Context, binding TicketBinding) error {
	s.mu.Lock()
	record, exists := s.active[binding]
	s.mu.Unlock()
	if !exists {
		return ErrTicketBindingInactive
	}
	if err := s.Authority.ReleaseConsumedWriter(ctx, record.claims); err != nil {
		s.mu.Lock()
		s.lifecycleFailures++
		s.mu.Unlock()
		return err
	}
	s.mu.Lock()
	s.releasedWriters++
	s.mu.Unlock()
	return nil
}

func (s *TicketSocketServer) close(ctx context.Context, binding TicketBinding) error {
	s.mu.Lock()
	record, exists := s.active[binding]
	if exists {
		delete(s.active, binding)
	}
	s.mu.Unlock()
	if !exists {
		return ErrTicketBindingInactive
	}
	if err := s.Authority.CloseConsumedStream(ctx, record.claims); err != nil {
		s.mu.Lock()
		s.lifecycleFailures++
		s.mu.Unlock()
		return err
	}
	s.mu.Lock()
	s.closedStreams++
	s.mu.Unlock()
	return nil
}

func (s *TicketSocketServer) heartbeat(binding TicketBinding) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	record, exists := s.active[binding]
	if !exists {
		return ErrTicketBindingInactive
	}
	record.lastSeen = s.now()
	s.active[binding] = record
	s.heartbeats++
	return nil
}

func (s *TicketSocketServer) recordEvent(ctx context.Context, binding TicketBinding, event SessionEvent) error {
	if err := validateSessionEvent(event); err != nil {
		return err
	}
	s.mu.Lock()
	record, exists := s.active[binding]
	s.mu.Unlock()
	if !exists {
		return ErrTicketBindingInactive
	}
	if err := s.Authority.RecordConsumedStreamEvent(ctx, record.claims, event); err != nil {
		s.mu.Lock()
		s.sessionEventFailures++
		s.mu.Unlock()
		return err
	}
	s.mu.Lock()
	s.sessionEvents++
	s.mu.Unlock()
	return nil
}

func (s *TicketSocketServer) Snapshot() TicketServiceSnapshot {
	s.mu.Lock()
	defer s.mu.Unlock()
	return TicketServiceSnapshot{
		ActiveBindings: len(s.active), MaxActiveBindings: s.maxActive(),
		Consumed: s.consumed, Heartbeats: s.heartbeats,
		ReleasedWriters: s.releasedWriters, ClosedStreams: s.closedStreams,
		ExpiredStreams: s.expiredStreams, LifecycleFailures: s.lifecycleFailures,
		SessionEvents: s.sessionEvents, SessionEventFailures: s.sessionEventFailures,
	}
}

func (s *TicketSocketServer) ExpireInactive(ctx context.Context) int {
	cutoff := s.now().Add(-s.activeTTL())
	s.mu.Lock()
	expired := make([]TicketClaims, 0)
	for binding, record := range s.active {
		if !record.lastSeen.After(cutoff) {
			delete(s.active, binding)
			expired = append(expired, record.claims)
		}
	}
	s.expiredStreams += uint64(len(expired))
	s.mu.Unlock()
	for _, claims := range expired {
		if err := s.Authority.CloseConsumedStream(ctx, claims); err != nil {
			s.mu.Lock()
			s.lifecycleFailures++
			s.mu.Unlock()
		}
	}
	return len(expired)
}

func (s *TicketSocketServer) runJanitor(stop <-chan struct{}) {
	interval := s.activeTTL() / 3
	if interval < time.Second {
		interval = time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			ctx, cancel := context.WithTimeout(context.Background(), defaultTicketSocketTTL)
			s.ExpireInactive(ctx)
			cancel()
		case <-stop:
			return
		}
	}
}

func (s *TicketSocketServer) now() time.Time {
	if s.Clock != nil {
		return s.Clock().UTC()
	}
	return time.Now().UTC()
}

func (s *TicketSocketServer) activeTTL() time.Duration {
	if s.ActiveTTL > 0 {
		return s.ActiveTTL
	}
	return defaultActiveBindingTTL
}

func (s *TicketSocketServer) maxActive() int {
	if s.MaxActiveClaims > 0 {
		return s.MaxActiveClaims
	}
	return defaultActiveBindings
}

func pointerToTicketStatus(status TicketServiceSnapshot) *TicketServiceSnapshot { return &status }

func (s *TicketSocketServer) writeError(connection net.Conn, err error) {
	code := "lifecycle_failed"
	switch {
	case errors.Is(err, ErrInvalidTicket):
		code = "invalid_ticket"
	case errors.Is(err, ErrTicketCapacity):
		code = "capacity_reached"
	case errors.Is(err, ErrTicketBindingInactive):
		code = "binding_inactive"
	}
	s.writeResponse(connection, ticketSocketResponse{Code: code, Error: err.Error()})
}

func (*TicketSocketServer) writeResponse(connection net.Conn, response ticketSocketResponse) {
	_ = json.NewEncoder(connection).Encode(response)
}

type TicketSocketClient struct {
	path    string
	timeout time.Duration
}

func ListenTicketSocket(path string) (net.Listener, error) {
	if path == "" || !filepath.IsAbs(path) {
		return nil, errors.New("direct stream ticket socket path must be absolute")
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o750); err != nil {
		return nil, err
	}
	if info, err := os.Lstat(path); err == nil {
		if info.Mode()&os.ModeSocket == 0 {
			return nil, errors.New("refusing to replace non-socket direct stream ticket path")
		}
		if err := os.Remove(path); err != nil {
			return nil, err
		}
	} else if !os.IsNotExist(err) {
		return nil, err
	}
	listener, err := net.Listen("unix", path)
	if err != nil {
		return nil, err
	}
	if err := os.Chmod(path, 0o660); err != nil {
		_ = listener.Close()
		return nil, err
	}
	return listener, nil
}

func NewTicketSocketClient(path string) *TicketSocketClient {
	if path == "" {
		path = DefaultTicketSocketPath
	}
	return &TicketSocketClient{path: path, timeout: defaultTicketSocketTTL}
}

func (c *TicketSocketClient) ConsumeTicket(ctx context.Context, ticket string, binding TicketBinding) (TicketClaims, error) {
	response, err := c.exchange(ctx, ticketSocketRequest{
		Action: "consume", Ticket: ticket, SessionID: binding.SessionID, Channel: binding.Channel,
	})
	if err != nil {
		return TicketClaims{}, err
	}
	if response.Claims == nil {
		return TicketClaims{}, errors.New("tera-api returned no direct stream claims")
	}
	return *response.Claims, nil
}

func (c *TicketSocketClient) ReleaseDirectWriter(ctx context.Context, binding TicketBinding) error {
	_, err := c.exchange(ctx, ticketSocketRequest{
		Action: "release_writer", SessionID: binding.SessionID, Channel: binding.Channel,
	})
	return err
}

func (c *TicketSocketClient) HeartbeatDirectStream(ctx context.Context, binding TicketBinding) error {
	_, err := c.exchange(ctx, ticketSocketRequest{
		Action: "heartbeat", SessionID: binding.SessionID, Channel: binding.Channel,
	})
	return err
}

func (c *TicketSocketClient) RecordDirectStreamEvent(ctx context.Context, binding TicketBinding, event SessionEvent) error {
	_, err := c.exchange(ctx, ticketSocketRequest{
		Action: "event", SessionID: binding.SessionID, Channel: binding.Channel, Event: &event,
	})
	return err
}

func (c *TicketSocketClient) CloseDirectStream(ctx context.Context, binding TicketBinding) error {
	_, err := c.exchange(ctx, ticketSocketRequest{
		Action: "close", SessionID: binding.SessionID, Channel: binding.Channel,
	})
	return err
}

func (c *TicketSocketClient) Status(ctx context.Context) (TicketServiceSnapshot, error) {
	response, err := c.exchange(ctx, ticketSocketRequest{Action: "status"})
	if err != nil {
		return TicketServiceSnapshot{}, err
	}
	if response.Status == nil {
		return TicketServiceSnapshot{}, errors.New("tera-api returned no direct stream ticket status")
	}
	return *response.Status, nil
}

func (c *TicketSocketClient) exchange(ctx context.Context, request ticketSocketRequest) (ticketSocketResponse, error) {
	dialer := net.Dialer{Timeout: c.timeout}
	connection, err := dialer.DialContext(ctx, "unix", c.path)
	if err != nil {
		return ticketSocketResponse{}, errors.New("tera-api direct stream ticket service is not reachable")
	}
	defer connection.Close()
	deadline := time.Now().Add(c.timeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	_ = connection.SetDeadline(deadline)
	if err := json.NewEncoder(connection).Encode(request); err != nil {
		return ticketSocketResponse{}, errors.New("direct stream ticket request failed")
	}
	var response ticketSocketResponse
	if err := json.NewDecoder(io.LimitReader(connection, maxTicketSocketBytes)).Decode(&response); err != nil {
		return ticketSocketResponse{}, errors.New("tera-api returned no direct stream ticket response")
	}
	if response.Error != "" {
		switch response.Code {
		case "invalid_ticket":
			return ticketSocketResponse{}, ErrInvalidTicket
		case "capacity_reached":
			return ticketSocketResponse{}, ErrTicketCapacity
		case "binding_inactive":
			return ticketSocketResponse{}, ErrTicketBindingInactive
		default:
			return ticketSocketResponse{}, errors.New(response.Error)
		}
	}
	return response, nil
}