teraWeb/backend/internal/directstream/websocket_test.go

package directstream

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

type websocketTransportFactory struct {
	peers chan net.Conn
}

func (f *websocketTransportFactory) Open(context.Context, string) (io.ReadWriteCloser, error) {
	transport, peer := net.Pipe()
	f.peers <- peer
	return transport, nil
}

type websocketAuthority struct {
	store        *TicketStore
	mu           sync.Mutex
	closed       []TicketBinding
	released     []TicketBinding
	heartbeats   int
	heartbeatErr error
	events       []SessionEvent
	eventErr     error
}

func (a *websocketAuthority) ConsumeTicket(_ context.Context, value string, binding TicketBinding) (TicketClaims, error) {
	return a.store.Consume(value, binding)
}

func (a *websocketAuthority) CloseDirectStream(_ context.Context, binding TicketBinding) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.closed = append(a.closed, binding)
	return nil
}

func (a *websocketAuthority) ReleaseDirectWriter(_ context.Context, binding TicketBinding) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.released = append(a.released, binding)
	return nil
}

func (a *websocketAuthority) HeartbeatDirectStream(_ context.Context, _ TicketBinding) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.heartbeats++
	return a.heartbeatErr
}

func (a *websocketAuthority) RecordDirectStreamEvent(_ context.Context, _ TicketBinding, event SessionEvent) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.events = append(a.events, event)
	return a.eventErr
}

func TestWebSocketHeartbeatFailureClosesConnectionAndUpdatesMetrics(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 1)}
	broker, _ := NewBroker(factory, BrokerOptions{})
	defer broker.Close()
	authority := &websocketAuthority{
		store:        NewTicketStore(TicketStoreOptions{}),
		heartbeatErr: errors.New("ticket service unavailable"),
	}
	metrics := &WebSocketMetrics{}
	server := httptest.NewServer(NewWebSocketHandler(broker, authority, authority, WebSocketOptions{
		PongWait: 200 * time.Millisecond, PingInterval: 20 * time.Millisecond,
		WriteTimeout: 100 * time.Millisecond, Metrics: metrics,
	}))
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH01-heartbeat", Channel: "CH01", Actor: "admin", Access: AccessViewer}
	ticket := issueWebSocketTicket(t, authority.store, claims)
	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	defer connection.Close()
	authenticateWebSocket(t, connection, ticket)
	_ = readStatusFrame(t, connection)
	_ = readWelcomeFrame(t, connection)
	peer := <-factory.peers
	defer peer.Close()
	_ = connection.SetReadDeadline(time.Now().Add(time.Second))
	if _, _, err := connection.ReadMessage(); err == nil {
		t.Fatal("expected failed lifecycle heartbeat to close WebSocket")
	}
	if !waitFor(time.Second, func() bool { return metrics.Snapshot().HeartbeatFailures == 1 }) {
		t.Fatalf("metrics = %#v", metrics.Snapshot())
	}
}

func TestWebSocketAuthenticatesTicketAndTransfersWriterBinary(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 1)}
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatal(err)
	}
	defer broker.Close()
	authority := &websocketAuthority{store: NewTicketStore(TicketStoreOptions{})}
	server := httptest.NewServer(NewWebSocketHandler(broker, authority, authority, WebSocketOptions{}))
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH01-ws1", Channel: "CH01", Actor: "admin", Access: AccessWriter}
	ticket := issueWebSocketTicket(t, authority.store, claims)
	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	defer connection.Close()
	authenticateWebSocket(t, connection, ticket)
	status := readStatusFrame(t, connection)
	if status.Type != "session.status" || status.Access != AccessWriter || !status.Writer {
		t.Fatalf("status = %#v", status)
	}
	welcome := readWelcomeFrame(t, connection)
	for _, expected := range []string{"Welcome to TERA CORA-X48 PRO K", "Connected from: 127.0.0.1", "Serial channel: CH01", "User: admin", "Access: writer", "Session: direct-CH01-ws1"} {
		if !strings.Contains(welcome, expected) {
			t.Fatalf("welcome = %q, missing %q", welcome, expected)
		}
	}
	peer := <-factory.peers
	defer peer.Close()

	tx := []byte{0x00, 0x3a, 0xff, '\n'}
	if err := connection.WriteMessage(websocket.BinaryMessage, tx); err != nil {
		t.Fatal(err)
	}
	peerTX := make([]byte, len(tx))
	if _, err := io.ReadFull(peer, peerTX); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(peerTX, tx) {
		t.Fatalf("tx = %v", peerTX)
	}

	rx := []byte{0xfe, 0x00, 0x7f}
	writeDone := make(chan error, 1)
	go func() {
		_, err := peer.Write(rx)
		writeDone <- err
	}()
	messageType, payload, err := connection.ReadMessage()
	if err != nil {
		t.Fatal(err)
	}
	if err := <-writeDone; err != nil {
		t.Fatal(err)
	}
	if messageType != websocket.BinaryMessage || !reflect.DeepEqual(payload, rx) {
		t.Fatalf("rx type=%d payload=%v", messageType, payload)
	}
}

func TestWebSocketConsumesTicketAndClosesLifecycleThroughAPISocket(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 1)}
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatal(err)
	}
	defer broker.Close()
	authority := &testTicketAuthority{store: NewTicketStore(TicketStoreOptions{})}
	ticketSocket, stopTickets := startTicketSocketServer(t, authority)
	defer stopTickets()
	ticketClient := NewTicketSocketClient(ticketSocket)
	server := httptest.NewServer(NewWebSocketHandler(broker, ticketClient, ticketClient, WebSocketOptions{}))
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH01-ipc", Channel: "CH01", Actor: "admin", Access: AccessViewer}
	ticket := issueWebSocketTicket(t, authority.store, claims)
	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	authenticateWebSocket(t, connection, ticket)
	status := readStatusFrame(t, connection)
	if status.Access != AccessViewer {
		t.Fatalf("status = %#v", status)
	}
	_ = readWelcomeFrame(t, connection)
	peer := <-factory.peers
	defer peer.Close()
	connection.Close()
	deadline := time.Now().Add(time.Second)
	for {
		authority.mu.Lock()
		closed := len(authority.closed)
		authority.mu.Unlock()
		if closed == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("ticket socket lifecycle close was not called")
		}
		time.Sleep(time.Millisecond)
	}
	authority.mu.Lock()
	events := append([]SessionEvent(nil), authority.events...)
	authority.mu.Unlock()
	if len(events) != 2 || events[0].Type != SessionEventConnected || events[1].Type != SessionEventDisconnected || events[0].Source != "127.0.0.1" {
		t.Fatalf("WebSocket session events = %#v", events)
	}
}

func TestWebSocketRejectsOriginSubprotocolReplayAndViewerWrites(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 2)}
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatal(err)
	}
	defer broker.Close()
	authority := &websocketAuthority{store: NewTicketStore(TicketStoreOptions{})}
	server := httptest.NewServer(NewWebSocketHandler(broker, authority, authority, WebSocketOptions{}))
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH01-ws2", Channel: "CH01", Actor: "viewer", Access: AccessViewer}
	ticket := issueWebSocketTicket(t, authority.store, claims)

	url := directWebSocketURL(server.URL, claims)
	badOriginHeader := http.Header{"Origin": []string{"https://evil.example"}}
	badOriginDialer := websocket.Dialer{Subprotocols: []string{Subprotocol}}
	if connection, response, err := badOriginDialer.Dial(url, badOriginHeader); err == nil || response == nil || response.StatusCode != http.StatusForbidden {
		if connection != nil {
			connection.Close()
		}
		t.Fatalf("bad Origin result response=%v err=%v", response, err)
	}
	noProtocolHeader := http.Header{"Origin": []string{server.URL}}
	if connection, response, err := websocket.DefaultDialer.Dial(url, noProtocolHeader); err == nil || response == nil || response.StatusCode != http.StatusUpgradeRequired {
		if connection != nil {
			connection.Close()
		}
		t.Fatalf("missing protocol result response=%v err=%v", response, err)
	}

	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	authenticateWebSocket(t, connection, ticket)
	_ = readStatusFrame(t, connection)
	_ = readWelcomeFrame(t, connection)
	peer := <-factory.peers
	defer peer.Close()
	if err := connection.WriteMessage(websocket.BinaryMessage, []byte{0x3a}); err != nil {
		t.Fatal(err)
	}
	errorFrame := readStatusFrame(t, connection)
	if errorFrame.Type != "error" || errorFrame.Code != "writer_required" {
		t.Fatalf("viewer error = %#v", errorFrame)
	}
	_ = peer.SetReadDeadline(time.Now().Add(50 * time.Millisecond))
	if _, err := peer.Read(make([]byte, 1)); err == nil {
		t.Fatal("viewer payload reached serial transport")
	}
	connection.Close()

	replay := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	defer replay.Close()
	authenticateWebSocket(t, replay, ticket)
	_, _, err = replay.ReadMessage()
	if err == nil {
		t.Fatal("expected replayed ticket connection to close")
	}
}

func TestWebSocketEnforcesFirstFrameDeadlineAndDataFrameLimit(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 1)}
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatal(err)
	}
	defer broker.Close()
	authority := &websocketAuthority{store: NewTicketStore(TicketStoreOptions{})}
	handler := NewWebSocketHandler(broker, authority, authority, WebSocketOptions{
		FirstFrameTimeout: 50 * time.Millisecond,
	})
	server := httptest.NewServer(handler)
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH01-ws3", Channel: "CH01", Actor: "admin", Access: AccessWriter}
	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	_ = connection.SetReadDeadline(time.Now().Add(time.Second))
	if _, _, err := connection.ReadMessage(); err == nil {
		t.Fatal("expected unauthenticated socket deadline")
	}
	connection.Close()

	claims.SessionID = "direct-CH01-ws4"
	ticket := issueWebSocketTicket(t, authority.store, claims)
	connection = dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	authenticateWebSocket(t, connection, ticket)
	_ = readStatusFrame(t, connection)
	_ = readWelcomeFrame(t, connection)
	peer := <-factory.peers
	defer peer.Close()
	oversized := make([]byte, MaxDataFrameBytes+1)
	if err := connection.WriteMessage(websocket.BinaryMessage, oversized); err != nil {
		t.Fatal(err)
	}
	_ = connection.SetReadDeadline(time.Now().Add(time.Second))
	if _, _, err := connection.ReadMessage(); err == nil {
		t.Fatal("expected oversized frame connection to close")
	}
}

func TestWebSocketRejectsConnectionWhenMandatorySessionLogFails(t *testing.T) {
	factory := &websocketTransportFactory{peers: make(chan net.Conn, 1)}
	broker, _ := NewBroker(factory, BrokerOptions{})
	defer broker.Close()
	authority := &websocketAuthority{
		store:    NewTicketStore(TicketStoreOptions{}),
		eventErr: errors.New("port log full"),
	}
	metrics := &WebSocketMetrics{}
	server := httptest.NewServer(NewWebSocketHandler(broker, authority, authority, WebSocketOptions{Metrics: metrics}))
	defer server.Close()
	claims := TicketClaims{SessionID: "direct-CH48-logfail", Channel: "CH48", Actor: "admin", Access: AccessViewer}
	ticket := issueWebSocketTicket(t, authority.store, claims)
	connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	defer connection.Close()
	authenticateWebSocket(t, connection, ticket)
	frame := readStatusFrame(t, connection)
	if frame.Type != "error" || frame.Code != "session_log_unavailable" {
		t.Fatalf("error frame = %#v", frame)
	}
	if !waitFor(time.Second, func() bool { return metrics.Snapshot().SessionLogFailures >= 1 }) {
		t.Fatalf("metrics = %#v", metrics.Snapshot())
	}
}

func TestDirectStreamSourceAndWelcomeCoverForwardedClientAndLastChannel(t *testing.T) {
	request := httptest.NewRequest(http.MethodGet, "/", nil)
	request.RemoteAddr = "127.0.0.1:42000"
	request.Header.Set("X-Forwarded-For", "198.51.100.24")
	if source := directStreamSource(request); source != "198.51.100.24" {
		t.Fatalf("source = %q", source)
	}
	request.Header.Set("X-Forwarded-For", "203.0.113.1, 198.51.100.24")
	if source := directStreamSource(request); source != "127.0.0.1" {
		t.Fatalf("multi-hop forwarded source = %q, want loopback fallback", source)
	}
	event := SessionEvent{Type: SessionEventConnected, Source: "198.51.100.24", Timestamp: time.Date(2026, 8, 13, 3, 4, 5, 0, time.UTC)}
	message := string(welcomeFrame(TicketClaims{SessionID: "direct-CH48-final", Channel: "CH48", Actor: "operator", Access: AccessViewer}, event))
	for _, expected := range []string{"Serial channel: CH48", "Connected from: 198.51.100.24", "Connected at: 2026-08-13T03:04:05Z"} {
		if !strings.Contains(message, expected) {
			t.Fatalf("CH48 welcome = %q, missing %q", message, expected)
		}
	}
}

func issueWebSocketTicket(t *testing.T, store *TicketStore, claims TicketClaims) string {
	t.Helper()
	issued, err := store.Issue(claims)
	if err != nil {
		t.Fatal(err)
	}
	return issued.Value
}

func dialDirectWebSocket(t *testing.T, serverURL string, claims TicketClaims, origin string, protocol bool) *websocket.Conn {
	t.Helper()
	dialer := websocket.Dialer{}
	if protocol {
		dialer.Subprotocols = []string{Subprotocol}
	}
	header := http.Header{"Origin": []string{origin}}
	connection, response, err := dialer.Dial(directWebSocketURL(serverURL, claims), header)
	if err != nil {
		status := 0
		if response != nil {
			status = response.StatusCode
		}
		t.Fatalf("dial websocket status=%d: %v", status, err)
	}
	return connection
}

func directWebSocketURL(serverURL string, claims TicketClaims) string {
	return "ws" + strings.TrimPrefix(serverURL, "http") + "/api/v1/direct/" + claims.Channel + "/sessions/" + claims.SessionID + "/stream"
}

func authenticateWebSocket(t *testing.T, connection *websocket.Conn, ticket string) {
	t.Helper()
	payload, err := json.Marshal(ControlFrame{Type: ControlAuthenticate, Ticket: ticket})
	if err != nil {
		t.Fatal(err)
	}
	if err := connection.WriteMessage(websocket.TextMessage, payload); err != nil {
		t.Fatal(err)
	}
}

func readStatusFrame(t *testing.T, connection *websocket.Conn) WebSocketStatusFrame {
	t.Helper()
	_ = connection.SetReadDeadline(time.Now().Add(time.Second))
	messageType, payload, err := connection.ReadMessage()
	if err != nil {
		var closeError *websocket.CloseError
		if errors.As(err, &closeError) {
			t.Fatalf("websocket closed code=%d text=%s", closeError.Code, closeError.Text)
		}
		t.Fatal(err)
	}
	if messageType != websocket.TextMessage {
		t.Fatalf("message type = %d, want text", messageType)
	}
	var frame WebSocketStatusFrame
	if err := json.Unmarshal(payload, &frame); err != nil {
		t.Fatal(err)
	}
	return frame
}

func readWelcomeFrame(t *testing.T, connection *websocket.Conn) string {
	t.Helper()
	_ = connection.SetReadDeadline(time.Now().Add(time.Second))
	messageType, payload, err := connection.ReadMessage()
	if err != nil {
		t.Fatal(err)
	}
	if messageType != websocket.BinaryMessage {
		t.Fatalf("welcome message type = %d, want binary", messageType)
	}
	return string(payload)
}