teraWeb/backend/internal/directstream/websocket.go

package directstream

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

const (
	defaultFirstFrameTimeout = 5 * time.Second
	defaultWebSocketPongWait = 30 * time.Second
	defaultWebSocketPing     = 10 * time.Second
	defaultWebSocketWrite    = 5 * time.Second
)

var directSessionIDPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$`)

type WebSocketOptions struct {
	AllowedOrigins    []string
	FirstFrameTimeout time.Duration
	PongWait          time.Duration
	PingInterval      time.Duration
	WriteTimeout      time.Duration
	Metrics           *WebSocketMetrics
}

type WebSocketSnapshot struct {
	ActiveConnections      int64  `json:"active_connections"`
	UpgradedConnections    uint64 `json:"upgraded_connections"`
	AcceptedConnections    uint64 `json:"accepted_connections"`
	ClosedConnections      uint64 `json:"closed_connections"`
	OriginRejections       uint64 `json:"origin_rejections"`
	ProtocolRejections     uint64 `json:"protocol_rejections"`
	AuthenticationFailures uint64 `json:"authentication_failures"`
	AttachmentFailures     uint64 `json:"attachment_failures"`
	LifecycleFailures      uint64 `json:"lifecycle_failures"`
	HeartbeatFailures      uint64 `json:"heartbeat_failures"`
	OutboundFailures       uint64 `json:"outbound_failures"`
	WelcomeFrames          uint64 `json:"welcome_frames"`
	SessionLogFailures     uint64 `json:"session_log_failures"`
}

type WebSocketMetrics struct {
	activeConnections      atomic.Int64
	upgradedConnections    atomic.Uint64
	acceptedConnections    atomic.Uint64
	closedConnections      atomic.Uint64
	originRejections       atomic.Uint64
	protocolRejections     atomic.Uint64
	authenticationFailures atomic.Uint64
	attachmentFailures     atomic.Uint64
	lifecycleFailures      atomic.Uint64
	heartbeatFailures      atomic.Uint64
	outboundFailures       atomic.Uint64
	welcomeFrames          atomic.Uint64
	sessionLogFailures     atomic.Uint64
}

func (m *WebSocketMetrics) Snapshot() WebSocketSnapshot {
	if m == nil {
		return WebSocketSnapshot{}
	}
	return WebSocketSnapshot{
		ActiveConnections:   m.activeConnections.Load(),
		UpgradedConnections: m.upgradedConnections.Load(),
		AcceptedConnections: m.acceptedConnections.Load(),
		ClosedConnections:   m.closedConnections.Load(),
		OriginRejections:    m.originRejections.Load(), ProtocolRejections: m.protocolRejections.Load(),
		AuthenticationFailures: m.authenticationFailures.Load(), AttachmentFailures: m.attachmentFailures.Load(),
		LifecycleFailures: m.lifecycleFailures.Load(), HeartbeatFailures: m.heartbeatFailures.Load(),
		OutboundFailures: m.outboundFailures.Load(), WelcomeFrames: m.welcomeFrames.Load(),
		SessionLogFailures: m.sessionLogFailures.Load(),
	}
}

type WebSocketStatusFrame struct {
	Type          string `json:"type"`
	Code          string `json:"code,omitempty"`
	Message       string `json:"message,omitempty"`
	SessionID     string `json:"session_id,omitempty"`
	Channel       string `json:"channel,omitempty"`
	Access        string `json:"access,omitempty"`
	Writer        bool   `json:"writer,omitempty"`
	QueuePosition int    `json:"queue_position,omitempty"`
}

type webSocketHandler struct {
	broker    *Broker
	tickets   TicketConsumer
	lifecycle TicketLifecycle
	options   WebSocketOptions
	origins   map[string]struct{}
	upgrader  websocket.Upgrader
	metrics   *WebSocketMetrics
}

func NewWebSocketHandler(broker *Broker, tickets TicketConsumer, lifecycle TicketLifecycle, options WebSocketOptions) http.Handler {
	if options.FirstFrameTimeout <= 0 {
		options.FirstFrameTimeout = defaultFirstFrameTimeout
	}
	if options.PongWait <= 0 {
		options.PongWait = defaultWebSocketPongWait
	}
	if options.PingInterval <= 0 || options.PingInterval >= options.PongWait {
		options.PingInterval = defaultWebSocketPing
	}
	if options.WriteTimeout <= 0 {
		options.WriteTimeout = defaultWebSocketWrite
	}
	origins := make(map[string]struct{}, len(options.AllowedOrigins))
	for _, origin := range options.AllowedOrigins {
		origin = strings.TrimSpace(origin)
		if origin != "" {
			origins[origin] = struct{}{}
		}
	}
	handler := &webSocketHandler{
		broker: broker, tickets: tickets, lifecycle: lifecycle,
		options: options, origins: origins, metrics: options.Metrics,
	}
	if handler.metrics == nil {
		handler.metrics = &WebSocketMetrics{}
	}
	handler.upgrader = websocket.Upgrader{
		ReadBufferSize:  4096,
		WriteBufferSize: 4096,
		Subprotocols:    []string{Subprotocol},
		CheckOrigin:     func(*http.Request) bool { return true },
	}
	return handler
}

func (h *webSocketHandler) ServeHTTP(w http.ResponseWriter, request *http.Request) {
	if h.broker == nil || h.tickets == nil || h.lifecycle == nil {
		http.Error(w, "Direct Mode streaming is unavailable", http.StatusServiceUnavailable)
		return
	}
	if request.Method != http.MethodGet {
		w.Header().Set("Allow", http.MethodGet)
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	channel, sessionID, ok := parseDirectStreamPath(request.URL.Path)
	if !ok {
		http.NotFound(w, request)
		return
	}
	if !h.originAllowed(request) {
		h.metrics.originRejections.Add(1)
		http.Error(w, "WebSocket Origin is not allowed", http.StatusForbidden)
		return
	}
	if !containsString(websocket.Subprotocols(request), Subprotocol) {
		h.metrics.protocolRejections.Add(1)
		w.Header().Set("Sec-WebSocket-Protocol", Subprotocol)
		http.Error(w, "WebSocket subprotocol is required", http.StatusUpgradeRequired)
		return
	}
	connection, err := h.upgrader.Upgrade(w, request, nil)
	if err != nil {
		return
	}
	h.metrics.upgradedConnections.Add(1)
	h.serveConnection(connection, TicketBinding{SessionID: sessionID, Channel: channel}, directStreamSource(request))
}

func (h *webSocketHandler) serveConnection(connection *websocket.Conn, binding TicketBinding, source string) {
	defer func() {
		h.metrics.closedConnections.Add(1)
		_ = connection.Close()
	}()
	connection.SetReadLimit(MaxDataFrameBytes)
	_ = connection.SetReadDeadline(time.Now().Add(h.options.FirstFrameTimeout))
	messageType, payload, err := connection.ReadMessage()
	if err != nil || messageType != websocket.TextMessage || len(payload) > MaxControlFrameBytes {
		h.metrics.authenticationFailures.Add(1)
		h.closeConnection(connection, websocket.ClosePolicyViolation, "authentication required")
		return
	}
	frame, err := DecodeControlFrame(payload)
	if err != nil || frame.Type != ControlAuthenticate {
		h.metrics.authenticationFailures.Add(1)
		h.closeConnection(connection, websocket.ClosePolicyViolation, "invalid authentication frame")
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), h.options.FirstFrameTimeout)
	claims, err := h.tickets.ConsumeTicket(ctx, frame.Ticket, binding)
	cancel()
	if err != nil {
		h.metrics.authenticationFailures.Add(1)
		h.closeConnection(connection, websocket.ClosePolicyViolation, "invalid stream ticket")
		return
	}
	attachment, err := h.broker.Attach(context.Background(), AttachRequest{
		SessionID: claims.SessionID, Channel: claims.Channel, Actor: claims.Actor, Access: claims.Access,
		SharedWrite: claims.Access == AccessWriter, MaxAttachments: claims.MaxAttachments,
	})
	if err != nil {
		h.metrics.attachmentFailures.Add(1)
		if h.closeLifecycle(binding) != nil {
			h.metrics.lifecycleFailures.Add(1)
		}
		reason := "stream attachment unavailable"
		if errors.Is(err, ErrAttachmentLimit) {
			reason = "channel session limit reached"
		}
		h.closeConnection(connection, websocket.CloseTryAgainLater, reason)
		return
	}
	h.metrics.acceptedConnections.Add(1)
	h.metrics.activeConnections.Add(1)
	defer h.metrics.activeConnections.Add(-1)
	defer attachment.Close()
	connectedLogged := false
	defer func() {
		if connectedLogged {
			event := SessionEvent{Type: SessionEventDisconnected, Source: source, Timestamp: time.Now().UTC()}
			if h.recordSessionEvent(binding, event) != nil {
				h.metrics.sessionLogFailures.Add(1)
			}
		}
		if h.closeLifecycle(binding) != nil {
			h.metrics.lifecycleFailures.Add(1)
		}
	}()

	_ = connection.SetReadDeadline(time.Now().Add(h.options.PongWait))
	connection.SetPongHandler(func(string) error {
		return connection.SetReadDeadline(time.Now().Add(h.options.PongWait))
	})
	var writeMu sync.Mutex
	writeJSON := func(value WebSocketStatusFrame) error {
		writeMu.Lock()
		defer writeMu.Unlock()
		_ = connection.SetWriteDeadline(time.Now().Add(h.options.WriteTimeout))
		return connection.WriteJSON(value)
	}
	writeBinary := func(value []byte) error {
		writeMu.Lock()
		defer writeMu.Unlock()
		_ = connection.SetWriteDeadline(time.Now().Add(h.options.WriteTimeout))
		return connection.WriteMessage(websocket.BinaryMessage, value)
	}
	writePing := func() error {
		writeMu.Lock()
		defer writeMu.Unlock()
		return connection.WriteControl(websocket.PingMessage, nil, time.Now().Add(h.options.WriteTimeout))
	}
	connectedEvent := SessionEvent{Type: SessionEventConnected, Source: source, Timestamp: time.Now().UTC()}
	if err := h.recordSessionEvent(binding, connectedEvent); err != nil {
		h.metrics.sessionLogFailures.Add(1)
		_ = writeJSON(errorStatus("session_log_unavailable", "Direct Mode session log is unavailable"))
		h.closeConnection(connection, websocket.CloseInternalServerErr, "session log unavailable")
		return
	}
	connectedLogged = true
	if err := writeJSON(h.statusFrame(claims.SessionID, claims.Channel, claims.Access)); err != nil {
		return
	}
	if err := writeBinary(welcomeFrame(claims, connectedEvent)); err != nil {
		h.metrics.outboundFailures.Add(1)
		return
	}
	h.metrics.welcomeFrames.Add(1)

	done := make(chan struct{})
	var stopOnce sync.Once
	stop := func() {
		stopOnce.Do(func() {
			close(done)
			_ = connection.Close()
		})
	}
	defer stop()
	go func() {
		for {
			data, err := attachment.Receive(context.Background())
			if err != nil {
				h.metrics.outboundFailures.Add(1)
				stop()
				return
			}
			if writeBinary(data) != nil {
				h.metrics.outboundFailures.Add(1)
				stop()
				return
			}
		}
	}()
	go func() {
		ticker := time.NewTicker(h.options.PingInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if writePing() != nil {
					h.metrics.outboundFailures.Add(1)
					stop()
					return
				}
				ctx, cancel := context.WithTimeout(context.Background(), h.options.WriteTimeout)
				err := h.lifecycle.HeartbeatDirectStream(ctx, binding)
				cancel()
				if err != nil {
					h.metrics.heartbeatFailures.Add(1)
					h.metrics.lifecycleFailures.Add(1)
					stop()
					return
				}
			case <-done:
				return
			}
		}
	}()

	access := claims.Access
	for {
		messageType, payload, err = connection.ReadMessage()
		if err != nil {
			return
		}
		switch messageType {
		case websocket.BinaryMessage:
			if err := ValidateClientData(access, true, payload); err != nil {
				if writeJSON(errorStatus("writer_required", err.Error())) != nil {
					return
				}
				continue
			}
			if err := attachment.Write(context.Background(), payload); err != nil {
				code := "stream_write_failed"
				if errors.Is(err, ErrWriterRequired) {
					code = "writer_required"
				}
				if writeJSON(errorStatus(code, err.Error())) != nil {
					return
				}
			}
		case websocket.TextMessage:
			if len(payload) > MaxControlFrameBytes {
				h.closeConnection(connection, websocket.CloseMessageTooBig, "control frame too large")
				return
			}
			control, err := DecodeControlFrame(payload)
			if err != nil || control.Type == ControlAuthenticate {
				h.closeConnection(connection, websocket.ClosePolicyViolation, "invalid control frame")
				return
			}
			switch control.Type {
			case ControlWriterRequest:
				if access != AccessWriter {
					if writeJSON(errorStatus("writer_required", "writer authorization requires a new REST claim")) != nil {
						return
					}
					continue
				}
				if _, err := attachment.RequestWriter(); err != nil {
					if writeJSON(errorStatus("writer_request_failed", err.Error())) != nil {
						return
					}
					continue
				}
			case ControlWriterRelease:
				if err := attachment.ReleaseWriter(); err != nil {
					if writeJSON(errorStatus("writer_release_failed", err.Error())) != nil {
						return
					}
					continue
				}
				if err := h.lifecycle.ReleaseDirectWriter(context.Background(), binding); err != nil {
					h.metrics.lifecycleFailures.Add(1)
					if writeJSON(errorStatus("writer_release_failed", err.Error())) != nil {
						return
					}
					continue
				}
				access = AccessViewer
			}
			if writeJSON(h.statusFrame(claims.SessionID, claims.Channel, access)) != nil {
				return
			}
		default:
			h.closeConnection(connection, websocket.CloseUnsupportedData, "unsupported frame type")
			return
		}
	}
}

func (h *webSocketHandler) statusFrame(sessionID, channel, access string) WebSocketStatusFrame {
	snapshot, _ := h.broker.Session(sessionID)
	return WebSocketStatusFrame{
		Type: "session.status", SessionID: sessionID, Channel: channel,
		Access: access, Writer: snapshot.Writer, QueuePosition: snapshot.QueuePosition,
	}
}

func (h *webSocketHandler) closeConnection(connection *websocket.Conn, code int, reason string) {
	deadline := time.Now().Add(h.options.WriteTimeout)
	_ = connection.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(code, reason), deadline)
}

func (h *webSocketHandler) closeLifecycle(binding TicketBinding) error {
	ctx, cancel := context.WithTimeout(context.Background(), h.options.WriteTimeout)
	defer cancel()
	return h.lifecycle.CloseDirectStream(ctx, binding)
}

func (h *webSocketHandler) recordSessionEvent(binding TicketBinding, event SessionEvent) error {
	ctx, cancel := context.WithTimeout(context.Background(), h.options.WriteTimeout)
	defer cancel()
	return h.lifecycle.RecordDirectStreamEvent(ctx, binding, event)
}

func directStreamSource(request *http.Request) string {
	host, _, err := net.SplitHostPort(strings.TrimSpace(request.RemoteAddr))
	if err != nil {
		host = strings.Trim(strings.TrimSpace(request.RemoteAddr), "[]")
	}
	remoteIP := net.ParseIP(host)
	if remoteIP == nil {
		return host
	}
	if remoteIP.IsLoopback() {
		forwarded := strings.Split(request.Header.Get("X-Forwarded-For"), ",")
		if len(forwarded) == 1 {
			if forwardedIP := net.ParseIP(strings.TrimSpace(forwarded[0])); forwardedIP != nil {
				return forwardedIP.String()
			}
		}
	}
	return remoteIP.String()
}

func (h *webSocketHandler) originAllowed(request *http.Request) bool {
	origin := request.Header.Get("Origin")
	if origin == "" {
		return false
	}
	if len(h.origins) != 0 {
		_, ok := h.origins[origin]
		return ok
	}
	parsed, err := url.Parse(origin)
	if err != nil || parsed.Host == "" || parsed.Host != request.Host {
		return false
	}
	expectedScheme := "http"
	if request.TLS != nil || strings.EqualFold(request.Header.Get("X-Forwarded-Proto"), "https") {
		expectedScheme = "https"
	}
	return parsed.Scheme == expectedScheme
}

func parseDirectStreamPath(path string) (string, string, bool) {
	parts := strings.Split(strings.Trim(path, "/"), "/")
	if len(parts) != 7 || parts[0] != "api" || parts[1] != "v1" || parts[2] != "direct" || parts[4] != "sessions" || parts[6] != "stream" {
		return "", "", false
	}
	if ValidateChannel(parts[3]) != nil || !directSessionIDPattern.MatchString(parts[5]) {
		return "", "", false
	}
	return parts[3], parts[5], true
}

func containsString(values []string, expected string) bool {
	for _, value := range values {
		if value == expected {
			return true
		}
	}
	return false
}

func errorStatus(code, message string) WebSocketStatusFrame {
	return WebSocketStatusFrame{Type: "error", Code: code, Message: message}
}