teraWeb/backend/internal/directstream/broker_test.go

package directstream

import (
	"bytes"
	"context"
	"errors"
	"io"
	"sync"
	"testing"
	"time"
)

func TestBrokerEnforcesFourAttachmentsAndClosesIdleTransport(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatalf("NewBroker() error = %v", err)
	}
	attachments := make([]*Attachment, 0, 4)
	for index := 1; index <= 4; index++ {
		attachment, err := broker.Attach(context.Background(), AttachRequest{
			SessionID: sessionID(index), Channel: "CH01", Actor: "viewer", Access: AccessViewer,
		})
		if err != nil {
			t.Fatalf("Attach(%d) error = %v", index, err)
		}
		attachments = append(attachments, attachment)
	}
	if _, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(5), Channel: "CH01", Actor: "viewer", Access: AccessViewer,
	}); !errors.Is(err, ErrAttachmentLimit) {
		t.Fatalf("Attach(fifth) error = %v, want ErrAttachmentLimit", err)
	}
	snapshot := broker.Snapshot()
	if snapshot.TotalAttachments != 4 || len(snapshot.Channels) != 1 || snapshot.Channels[0].AttachmentCount != 4 {
		t.Fatalf("Snapshot() = %#v", snapshot)
	}
	transport := factory.transport("CH01")
	for _, attachment := range attachments {
		if err := attachment.Close(); err != nil {
			t.Fatalf("Close() error = %v", err)
		}
	}
	if !waitFor(time.Second, func() bool { return transport.isClosed() }) {
		t.Fatal("idle channel transport was not closed")
	}
	if got := broker.Snapshot().TotalAttachments; got != 0 {
		t.Fatalf("TotalAttachments = %d, want 0", got)
	}
}

func TestBrokerAllowsOnlyWriterToSendBinaryData(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	writer, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(1), Channel: "CH01", Actor: "operator", Access: AccessWriter,
	})
	if err != nil {
		t.Fatalf("Attach(writer) error = %v", err)
	}
	viewer, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(2), Channel: "CH01", Actor: "viewer", Access: AccessViewer,
	})
	if err != nil {
		t.Fatalf("Attach(viewer) error = %v", err)
	}
	if err := viewer.Write(context.Background(), []byte{0x3a}); !errors.Is(err, ErrWriterRequired) {
		t.Fatalf("viewer.Write() error = %v, want ErrWriterRequired", err)
	}
	payload := []byte{0x00, 0x3a, 0xff, 0x0d, 0x0a}
	if err := writer.Write(context.Background(), payload); err != nil {
		t.Fatalf("writer.Write() error = %v", err)
	}
	if got := factory.transport("CH01").written(); !bytes.Equal(got, payload) {
		t.Fatalf("transport bytes = %v, want %v", got, payload)
	}
}

func TestBrokerSharedWritersStayConnectedAndSerializeUARTWrites(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	first, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(1), Channel: "CH01", Actor: "terminal-1", Access: AccessWriter, SharedWrite: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	second, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(2), Channel: "CH01", Actor: "terminal-2", Access: AccessWriter, SharedWrite: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()
	defer second.Close()

	if err := first.Write(context.Background(), []byte("first\r")); err != nil {
		t.Fatalf("first shared write: %v", err)
	}
	if err := second.Write(context.Background(), []byte("second\r")); err != nil {
		t.Fatalf("second shared write: %v", err)
	}
	if got := string(factory.transport("CH01").written()); got != "first\rsecond\r" {
		t.Fatalf("serialized UART writes = %q", got)
	}
	for _, attachment := range []*Attachment{first, second} {
		if snapshot, ok := broker.Session(attachment.SessionID()); !ok || !snapshot.Writer || snapshot.QueuePosition != 0 {
			t.Fatalf("shared session snapshot = %#v, %v", snapshot, ok)
		}
	}
}

func TestBrokerSharedWriterCanSendSerialBreak(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	attachment, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(1), Channel: "CH01", Actor: "terminal", Access: AccessWriter, SharedWrite: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer attachment.Close()
	if err := attachment.SendBreak(context.Background()); err != nil {
		t.Fatalf("SendBreak() error = %v", err)
	}
	if got := factory.transport("CH01").breakCount(); got != 1 {
		t.Fatalf("break count = %d, want 1", got)
	}
}

func TestBrokerWriterFIFOReleaseAndDisconnect(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	first := mustAttach(t, broker, 1, AccessWriter)
	second := mustAttach(t, broker, 2, AccessViewer)
	third := mustAttach(t, broker, 3, AccessViewer)

	if position, err := second.RequestWriter(); err != nil || position != 1 {
		t.Fatalf("second.RequestWriter() = (%d, %v), want (1, nil)", position, err)
	}
	if position, err := third.RequestWriter(); err != nil || position != 2 {
		t.Fatalf("third.RequestWriter() = (%d, %v), want (2, nil)", position, err)
	}
	if err := first.ReleaseWriter(); err != nil {
		t.Fatalf("first.ReleaseWriter() error = %v", err)
	}
	if snapshot, _ := broker.Session(second.SessionID()); !snapshot.Writer {
		t.Fatalf("second session = %#v, want writer", snapshot)
	}
	if err := second.Close(); err != nil {
		t.Fatalf("second.Close() error = %v", err)
	}
	if snapshot, _ := broker.Session(third.SessionID()); !snapshot.Writer {
		t.Fatalf("third session = %#v, want writer after disconnect", snapshot)
	}
}

func TestBrokerFansOutRXBytes(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	first := mustAttach(t, broker, 1, AccessViewer)
	second := mustAttach(t, broker, 2, AccessViewer)
	payload := []byte{0x00, 0x3a, 0xff}
	factory.transport("CH01").emit(payload)

	if got := receiveWithTimeout(t, first); !bytes.Equal(got, payload) {
		t.Fatalf("first.Receive() = %v, want %v", got, payload)
	}
	if got := receiveWithTimeout(t, second); !bytes.Equal(got, payload) {
		t.Fatalf("second.Receive() = %v, want %v", got, payload)
	}
}

func TestBrokerDisconnectsOnlySlowAttachment(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{AttachmentQueueBytes: 4})
	slow := mustAttach(t, broker, 1, AccessViewer)
	fast := mustAttach(t, broker, 2, AccessViewer)
	transport := factory.transport("CH01")
	transport.emit([]byte{1, 2, 3, 4})
	if got := receiveWithTimeout(t, fast); !bytes.Equal(got, []byte{1, 2, 3, 4}) {
		t.Fatalf("fast first receive = %v", got)
	}
	transport.emit([]byte{5, 6})
	if !waitFor(time.Second, func() bool { return broker.Snapshot().TotalAttachments == 1 }) {
		t.Fatalf("slow attachment was not isolated: %#v", broker.Snapshot())
	}
	if got := receiveWithTimeout(t, fast); !bytes.Equal(got, []byte{5, 6}) {
		t.Fatalf("fast second receive = %v", got)
	}
	if got := receiveWithTimeout(t, slow); !bytes.Equal(got, []byte{1, 2, 3, 4}) {
		t.Fatalf("slow queued receive = %v", got)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if _, err := slow.Receive(ctx); !errors.Is(err, ErrSlowConsumer) {
		t.Fatalf("slow.Receive() error = %v, want ErrSlowConsumer", err)
	}
	if got := broker.Snapshot().SlowConsumerDisconnects; got != 1 {
		t.Fatalf("SlowConsumerDisconnects = %d, want 1", got)
	}
}

func TestBrokerSlowOpenOnOneChannelDoesNotBlockExistingChannel(t *testing.T) {
	factory := newGatedTransportFactory("CH02")
	broker, err := NewBroker(factory, BrokerOptions{})
	if err != nil {
		t.Fatal(err)
	}
	first, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(1), Channel: "CH01", Actor: "operator", Access: AccessWriter,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer first.Close()

	opened := make(chan error, 1)
	go func() {
		_, attachErr := broker.Attach(context.Background(), AttachRequest{
			SessionID: sessionID(2), Channel: "CH02", Actor: "viewer", Access: AccessViewer,
		})
		opened <- attachErr
	}()
	select {
	case <-factory.started:
	case <-time.After(time.Second):
		t.Fatal("CH02 open did not start")
	}

	writable := make(chan error, 1)
	go func() { writable <- first.Write(context.Background(), []byte{0x3a}) }()
	select {
	case err := <-writable:
		if err != nil {
			t.Fatalf("existing CH01 write failed: %v", err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("CH02 device open blocked the existing CH01 data path")
	}
	close(factory.release)
	if err := <-opened; err != nil {
		t.Fatalf("CH02 attach failed: %v", err)
	}
}

func TestBrokerSnapshotRecordsTrafficFailureAndReconnect(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	first := mustAttach(t, broker, 1, AccessWriter)
	if err := first.Write(context.Background(), []byte{0x00, 0x3a, 0xff}); err != nil {
		t.Fatal(err)
	}
	transport := factory.transport("CH01")
	transport.emit([]byte{0x10, 0x11})
	_ = receiveWithTimeout(t, first)
	if err := first.Close(); err != nil {
		t.Fatal(err)
	}

	second := mustAttach(t, broker, 2, AccessViewer)
	secondTransport := factory.transport("CH01")
	if err := secondTransport.Close(); err != nil {
		t.Fatal(err)
	}
	if !waitFor(time.Second, func() bool { return broker.Snapshot().TotalAttachments == 0 }) {
		t.Fatal("serial disappearance did not detach the active session")
	}
	_ = second.Close()

	snapshot := broker.Snapshot()
	if snapshot.TXBytes != 3 || snapshot.RXBytes != 2 {
		t.Fatalf("traffic snapshot = %#v", snapshot)
	}
	if snapshot.TransportOpens != 2 || snapshot.TransportReconnects != 1 || snapshot.TransportFailures != 1 {
		t.Fatalf("transport snapshot = %#v", snapshot)
	}
}

func TestBrokerShutdownClosesAttachmentsWithVisibleBoundedFailure(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{})
	attachment := mustAttach(t, broker, 1, AccessViewer)
	if err := broker.Close(); err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if _, err := attachment.Receive(ctx); !errors.Is(err, ErrChannelClosed) {
		t.Fatalf("Receive() after broker shutdown = %v, want ErrChannelClosed", err)
	}
	if _, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(2), Channel: "CH01", Actor: "viewer", Access: AccessViewer,
	}); !errors.Is(err, ErrChannelClosed) {
		t.Fatalf("Attach() after broker shutdown = %v, want ErrChannelClosed", err)
	}
	if snapshot := broker.Snapshot(); snapshot.TotalAttachments != 0 || snapshot.RejectedAttachments != 1 {
		t.Fatalf("shutdown snapshot = %#v", snapshot)
	}
}

func TestBrokerConcurrentWriterAndRXPathsRemainBounded(t *testing.T) {
	factory := newFakeTransportFactory()
	broker, _ := NewBroker(factory, BrokerOptions{AttachmentQueueBytes: 64 * 1024})
	writer := mustAttach(t, broker, 1, AccessWriter)
	viewer := mustAttach(t, broker, 2, AccessViewer)
	transport := factory.transport("CH01")
	const iterations = 100
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		for index := 0; index < iterations; index++ {
			if err := writer.Write(context.Background(), []byte{byte(index)}); err != nil {
				t.Errorf("Write(%d) error = %v", index, err)
				return
			}
		}
	}()
	go func() {
		defer wg.Done()
		for index := 0; index < iterations; index++ {
			transport.emit([]byte{byte(index)})
		}
	}()
	for index := 0; index < iterations; index++ {
		receiveWithTimeout(t, viewer)
	}
	wg.Wait()
	if snapshot := broker.Snapshot(); snapshot.TotalAttachments != 2 || snapshot.Channels[0].QueuedBytes > 64*1024*2 {
		t.Fatalf("Snapshot() = %#v", snapshot)
	}
}

func mustAttach(t *testing.T, broker *Broker, index int, access string) *Attachment {
	t.Helper()
	attachment, err := broker.Attach(context.Background(), AttachRequest{
		SessionID: sessionID(index), Channel: "CH01", Actor: "actor", Access: access,
	})
	if err != nil {
		t.Fatalf("Attach(%d) error = %v", index, err)
	}
	return attachment
}

func sessionID(index int) string {
	return "direct_CH01_00000" + string(rune('0'+index))
}

func receiveWithTimeout(t *testing.T, attachment *Attachment) []byte {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	payload, err := attachment.Receive(ctx)
	if err != nil {
		t.Fatalf("Receive() error = %v", err)
	}
	return payload
}

func waitFor(timeout time.Duration, condition func() bool) bool {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if condition() {
			return true
		}
		time.Sleep(time.Millisecond)
	}
	return condition()
}

type fakeTransportFactory struct {
	mu         sync.Mutex
	transports map[string]*fakeTransport
}

type gatedTransportFactory struct {
	*fakeTransportFactory
	channel string
	started chan struct{}
	release chan struct{}
	once    sync.Once
}

func newGatedTransportFactory(channel string) *gatedTransportFactory {
	return &gatedTransportFactory{
		fakeTransportFactory: newFakeTransportFactory(), channel: channel,
		started: make(chan struct{}), release: make(chan struct{}),
	}
}

func (f *gatedTransportFactory) Open(ctx context.Context, channel string) (io.ReadWriteCloser, error) {
	if channel == f.channel {
		f.once.Do(func() { close(f.started) })
		select {
		case <-f.release:
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
	return f.fakeTransportFactory.Open(ctx, channel)
}

func newFakeTransportFactory() *fakeTransportFactory {
	return &fakeTransportFactory{transports: make(map[string]*fakeTransport)}
}

func (f *fakeTransportFactory) Open(_ context.Context, channel string) (io.ReadWriteCloser, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	transport := newFakeTransport()
	f.transports[channel] = transport
	return transport, nil
}

func (f *fakeTransportFactory) transport(channel string) *fakeTransport {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.transports[channel]
}

type fakeTransport struct {
	rx        chan []byte
	closed    chan struct{}
	closeOnce sync.Once
	mu        sync.Mutex
	writeData []byte
	pending   []byte
	breaks    int
}

func (t *fakeTransport) SendBreak(context.Context) error {
	t.mu.Lock()
	t.breaks++
	t.mu.Unlock()
	return nil
}

func newFakeTransport() *fakeTransport {
	return &fakeTransport{rx: make(chan []byte, 256), closed: make(chan struct{})}
}

func (t *fakeTransport) Read(buffer []byte) (int, error) {
	for len(t.pending) == 0 {
		select {
		case payload := <-t.rx:
			t.pending = payload
		case <-t.closed:
			return 0, io.EOF
		}
	}
	count := copy(buffer, t.pending)
	t.pending = t.pending[count:]
	return count, nil
}

func (t *fakeTransport) Write(payload []byte) (int, error) {
	select {
	case <-t.closed:
		return 0, io.ErrClosedPipe
	default:
	}
	t.mu.Lock()
	t.writeData = append(t.writeData, payload...)
	t.mu.Unlock()
	return len(payload), nil
}

func (t *fakeTransport) Close() error {
	t.closeOnce.Do(func() { close(t.closed) })
	return nil
}

func (t *fakeTransport) emit(payload []byte) {
	copyOfPayload := append([]byte(nil), payload...)
	select {
	case t.rx <- copyOfPayload:
	case <-t.closed:
	}
}

func (t *fakeTransport) written() []byte {
	t.mu.Lock()
	defer t.mu.Unlock()
	return append([]byte(nil), t.writeData...)
}

func (t *fakeTransport) breakCount() int {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.breaks
}

func (t *fakeTransport) isClosed() bool {
	select {
	case <-t.closed:
		return true
	default:
		return false
	}
}