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 } }