teraWeb/backend/internal/directstream/shared_socket_test.go

package directstream

import (
	"bytes"
	"context"
	"errors"
	"net/http/httptest"
	"os"
	"path/filepath"
	"testing"
	"time"

	"github.com/gorilla/websocket"
)

func sharedTestSocket(t *testing.T) string {
	t.Helper()
	dir, err := os.MkdirTemp("", "shared-")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = os.RemoveAll(dir) })
	return filepath.Join(dir, "s.sock")
}

func TestBrowserAndTerminalSessionsShareOneUARTAndCombinedLimit(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	defer broker.Close()
	path := sharedTestSocket(t)
	listener, err := ListenTicketSocket(path)
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	go ServeSharedSocket(listener, broker)
	client := SharedSocketClient{Path: path}
	authority := &websocketAuthority{store: NewTicketStore(TicketStoreOptions{})}
	server := httptest.NewServer(NewWebSocketHandler(broker, authority, authority, WebSocketOptions{}))
	defer server.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var browsers []*websocket.Conn
	for _, id := range []string{"browser-1", "browser-2"} {
		claims := TicketClaims{SessionID: id, Channel: "CH20", Actor: id, Access: AccessWriter, MaxAttachments: 4}
		connection := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
		defer connection.Close()
		authenticateWebSocket(t, connection, issueWebSocketTicket(t, authority.store, claims))
		status := readStatusFrame(t, connection)
		if !status.Writer || status.QueuePosition != 0 {
			t.Fatalf("browser waiting: %+v", status)
		}
		_ = readWelcomeFrame(t, connection)
		browsers = append(browsers, connection)
	}
	var terminals []SharedAttachment
	for _, id := range []string{"ssh-operation-1", "telnet-operation-2"} {
		attachment, err := client.Attach(ctx, AttachRequest{SessionID: id, Channel: "CH20", Actor: id, Access: AccessWriter, MaxAttachments: 4})
		if err != nil {
			t.Fatal(err)
		}
		defer attachment.Close()
		terminals = append(terminals, attachment)
	}
	if _, err := client.Attach(ctx, AttachRequest{SessionID: "fifth", Channel: "CH20", Actor: "fifth", Access: AccessWriter, MaxAttachments: 4}); !errors.Is(err, ErrAttachmentLimit) {
		t.Fatalf("fifth terminal: %v", err)
	}
	claims := TicketClaims{SessionID: "fifth-browser", Channel: "CH20", Actor: "fifth", Access: AccessWriter, MaxAttachments: 4}
	fifth := dialDirectWebSocket(t, server.URL, claims, server.URL, true)
	defer fifth.Close()
	authenticateWebSocket(t, fifth, issueWebSocketTicket(t, authority.store, claims))
	_ = fifth.SetReadDeadline(time.Now().Add(time.Second))
	if _, _, err := fifth.ReadMessage(); err == nil {
		t.Fatal("fifth browser accepted")
	}
	transport := factory.transport("CH20")
	for _, browser := range browsers {
		if err := browser.WriteMessage(websocket.BinaryMessage, []byte("web\r")); err != nil {
			t.Fatal(err)
		}
	}
	for _, terminal := range terminals {
		if err := terminal.Write(ctx, []byte("terminal\r")); err != nil {
			t.Fatal(err)
		}
	}
	if !waitFor(time.Second, func() bool { return len(transport.written()) == 26 }) {
		t.Fatalf("UART writes: %q", transport.written())
	}
	payload := []byte{0, 0xff, 'o', 'k', '\r', '\n'}
	transport.emit(payload)
	for _, browser := range browsers {
		_ = browser.SetReadDeadline(time.Now().Add(time.Second))
		kind, actual, err := browser.ReadMessage()
		if err != nil || kind != websocket.BinaryMessage || !bytes.Equal(actual, payload) {
			t.Fatalf("browser RX: %v %v", actual, err)
		}
	}
	for _, terminal := range terminals {
		actual, err := terminal.Receive(ctx)
		if err != nil || !bytes.Equal(actual, payload) {
			t.Fatalf("terminal RX: %v %v", actual, err)
		}
	}
	if err := terminals[0].SendBreak(ctx); err != nil || transport.breakCount() != 1 {
		t.Fatalf("break: %v", err)
	}
	_ = terminals[0].Close()
	if !waitFor(time.Second, func() bool { return broker.Snapshot().TotalAttachments == 3 }) {
		t.Fatal("closed terminal slot leaked")
	}
	replacement, err := client.Attach(ctx, AttachRequest{SessionID: "replacement", Channel: "CH20", Actor: "replacement", Access: AccessWriter, MaxAttachments: 4})
	if err != nil {
		t.Fatal(err)
	}
	defer replacement.Close()
	if err := terminals[1].Write(ctx, []byte("still-connected\r")); err != nil {
		t.Fatal(err)
	}
	if snapshot := broker.Snapshot(); snapshot.TransportOpens != 1 || snapshot.TotalAttachments != 4 {
		t.Fatalf("snapshot: %+v", snapshot)
	}
}

func TestSharedSocketHonorsLowerConfiguredLimitAndViewerAccess(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	defer broker.Close()
	path := sharedTestSocket(t)
	listener, err := ListenTicketSocket(path)
	if err != nil {
		t.Fatal(err)
	}
	defer listener.Close()
	go ServeSharedSocket(listener, broker)
	client := SharedSocketClient{Path: path}
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	viewer, err := client.Attach(ctx, AttachRequest{SessionID: "viewer", Channel: "CH01", Actor: "viewer", Access: AccessViewer, SharedWrite: true, MaxAttachments: 1})
	if err != nil {
		t.Fatal(err)
	}
	defer viewer.Close()
	if err := viewer.Write(ctx, []byte("denied")); err == nil {
		t.Fatal("viewer wrote data")
	}
	if _, err := client.Attach(ctx, AttachRequest{SessionID: "writer", Channel: "CH01", Actor: "writer", Access: AccessWriter, MaxAttachments: 4}); !errors.Is(err, ErrAttachmentLimit) {
		t.Fatalf("configured limit bypassed: %v", err)
	}
}