teraWeb/backend/internal/directstream/broker.go

package directstream

import (
	"context"
	"errors"
	"io"
	"sort"
	"sync"
	"sync/atomic"
)

const (
	defaultMaxAttachmentsPerChannel = 4
	defaultAttachmentQueueBytes     = 64 * 1024
	serialReadBufferBytes           = 32 * 1024
)

var (
	ErrAttachmentLimit = errors.New("direct stream attachment limit reached")
	ErrSessionExists   = errors.New("direct stream session already attached")
	ErrSessionNotFound = errors.New("direct stream session not found")
	ErrSlowConsumer    = errors.New("direct stream attachment is too slow")
	ErrChannelClosed   = errors.New("direct stream channel transport closed")
)

type TransportFactory interface {
	Open(ctx context.Context, channel string) (io.ReadWriteCloser, error)
}

type breakTransport interface {
	SendBreak(context.Context) error
}

type BrokerOptions struct {
	MaxAttachmentsPerChannel int
	AttachmentQueueBytes     int
}

type AttachRequest struct {
	SessionID string
	Channel   string
	Actor     string
	Access    string
	// SharedWrite is reserved for authenticated reverse SSH/Telnet services.
	// It lets every accepted terminal session write to the same UART while the
	// channel write mutex preserves each frame as one serialized operation.
	// Browser direct-mode sessions keep the existing single-writer FIFO policy.
	SharedWrite bool
}

type SessionSnapshot struct {
	SessionID     string `json:"session_id"`
	Channel       string `json:"channel"`
	Actor         string `json:"actor"`
	Writer        bool   `json:"writer"`
	QueuePosition int    `json:"queue_position"`
	OutboundBytes int    `json:"outbound_bytes"`
}

type ChannelSnapshot struct {
	Channel          string   `json:"channel"`
	AttachmentCount  int      `json:"attachment_count"`
	MaxAttachments   int      `json:"max_attachments"`
	WriterSessionID  string   `json:"writer_session_id,omitempty"`
	PendingWriterIDs []string `json:"pending_writer_ids,omitempty"`
	QueuedBytes      int      `json:"queued_bytes"`
}

type BrokerSnapshot struct {
	TotalAttachments        int               `json:"total_attachments"`
	ActiveChannels          int               `json:"active_channels"`
	AttachAttempts          uint64            `json:"attach_attempts"`
	SuccessfulAttachments   uint64            `json:"successful_attachments"`
	RejectedAttachments     uint64            `json:"rejected_attachments"`
	TransportOpens          uint64            `json:"transport_opens"`
	TransportOpenFailures   uint64            `json:"transport_open_failures"`
	TransportReconnects     uint64            `json:"transport_reconnects"`
	TransportFailures       uint64            `json:"transport_failures"`
	SlowConsumerDisconnects uint64            `json:"slow_consumer_disconnects"`
	RXBytes                 uint64            `json:"rx_bytes"`
	TXBytes                 uint64            `json:"tx_bytes"`
	FanoutBytes             uint64            `json:"fanout_bytes"`
	Channels                []ChannelSnapshot `json:"channels"`
}

type Broker struct {
	mu                  sync.Mutex
	factory             TransportFactory
	maxAttachments      int
	attachmentQueueSize int
	channels            map[string]*channelState
	sessions            map[string]*brokerSession
	openings            map[string]*channelOpening
	seenChannels        map[string]bool
	lifetime            context.Context
	cancel              context.CancelFunc
	metrics             brokerMetrics
	closed              bool
}

type channelState struct {
	channel       string
	transport     io.ReadWriteCloser
	sessions      map[string]*brokerSession
	writerID      string
	writerQueue   []string
	writeMu       sync.Mutex
	readerStarted bool
}

type channelOpening struct {
	done chan struct{}
	err  error
}

type brokerMetrics struct {
	attachAttempts          atomic.Uint64
	successfulAttachments   atomic.Uint64
	rejectedAttachments     atomic.Uint64
	transportOpens          atomic.Uint64
	transportOpenFailures   atomic.Uint64
	transportReconnects     atomic.Uint64
	transportFailures       atomic.Uint64
	slowConsumerDisconnects atomic.Uint64
	rxBytes                 atomic.Uint64
	txBytes                 atomic.Uint64
	fanoutBytes             atomic.Uint64
}

type brokerSession struct {
	request  AttachRequest
	outbound *byteQueue
}

type Attachment struct {
	broker    *Broker
	sessionID string
	outbound  *byteQueue
	closeOnce sync.Once
	closeErr  error
}

func NewBroker(factory TransportFactory, options BrokerOptions) (*Broker, error) {
	if factory == nil {
		return nil, errors.New("direct stream transport factory is required")
	}
	maxAttachments := options.MaxAttachmentsPerChannel
	if maxAttachments <= 0 {
		maxAttachments = defaultMaxAttachmentsPerChannel
	}
	queueBytes := options.AttachmentQueueBytes
	if queueBytes <= 0 {
		queueBytes = defaultAttachmentQueueBytes
	}
	lifetime, cancel := context.WithCancel(context.Background())
	return &Broker{
		factory:             factory,
		maxAttachments:      maxAttachments,
		attachmentQueueSize: queueBytes,
		channels:            make(map[string]*channelState),
		sessions:            make(map[string]*brokerSession),
		openings:            make(map[string]*channelOpening),
		seenChannels:        make(map[string]bool),
		lifetime:            lifetime,
		cancel:              cancel,
	}, nil
}

func (b *Broker) Attach(ctx context.Context, request AttachRequest) (*Attachment, error) {
	if request.SessionID == "" || request.Actor == "" {
		return nil, ErrInvalidTicketClaims
	}
	if err := ValidateChannel(request.Channel); err != nil {
		return nil, err
	}
	if err := ValidateAccess(request.Access); err != nil {
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	b.metrics.attachAttempts.Add(1)
	for {
		b.mu.Lock()
		if b.closed {
			b.mu.Unlock()
			b.metrics.rejectedAttachments.Add(1)
			return nil, ErrChannelClosed
		}
		if _, exists := b.sessions[request.SessionID]; exists {
			b.mu.Unlock()
			b.metrics.rejectedAttachments.Add(1)
			return nil, ErrSessionExists
		}
		if state := b.channels[request.Channel]; state != nil {
			if len(state.sessions) >= b.maxAttachments {
				b.mu.Unlock()
				b.metrics.rejectedAttachments.Add(1)
				return nil, ErrAttachmentLimit
			}
			session := &brokerSession{request: request, outbound: newByteQueue(b.attachmentQueueSize)}
			state.sessions[request.SessionID] = session
			b.sessions[request.SessionID] = session
			if request.Access == AccessWriter && !request.SharedWrite {
				if state.writerID == "" {
					state.writerID = request.SessionID
				} else {
					state.writerQueue = append(state.writerQueue, request.SessionID)
				}
			}
			startReader := !state.readerStarted
			state.readerStarted = true
			attachment := &Attachment{broker: b, sessionID: request.SessionID, outbound: session.outbound}
			b.mu.Unlock()
			b.metrics.successfulAttachments.Add(1)
			if startReader {
				go b.readChannel(state)
			}
			return attachment, nil
		}
		if opening := b.openings[request.Channel]; opening != nil {
			b.mu.Unlock()
			select {
			case <-ctx.Done():
				b.metrics.rejectedAttachments.Add(1)
				return nil, ctx.Err()
			case <-opening.done:
				if opening.err != nil {
					b.metrics.rejectedAttachments.Add(1)
					return nil, opening.err
				}
				continue
			}
		}

		opening := &channelOpening{done: make(chan struct{})}
		b.openings[request.Channel] = opening
		b.mu.Unlock()

		openContext, cancel := context.WithCancel(ctx)
		stopLifetimeCancel := context.AfterFunc(b.lifetime, cancel)
		transport, openErr := b.factory.Open(openContext, request.Channel)
		stopLifetimeCancel()
		cancel()

		b.mu.Lock()
		delete(b.openings, request.Channel)
		if openErr == nil && b.closed {
			openErr = ErrChannelClosed
		}
		if openErr == nil {
			b.channels[request.Channel] = &channelState{
				channel: request.Channel, transport: transport,
				sessions: make(map[string]*brokerSession),
			}
			b.metrics.transportOpens.Add(1)
			if b.seenChannels[request.Channel] {
				b.metrics.transportReconnects.Add(1)
			}
			b.seenChannels[request.Channel] = true
		} else {
			b.metrics.transportOpenFailures.Add(1)
		}
		opening.err = openErr
		close(opening.done)
		b.mu.Unlock()
		if openErr != nil {
			if transport != nil {
				_ = transport.Close()
			}
			b.metrics.rejectedAttachments.Add(1)
			return nil, openErr
		}
	}
}

func (b *Broker) RequestWriter(sessionID string) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		return 0, ErrSessionNotFound
	}
	if state.writerID == sessionID {
		return 0, nil
	}
	if position := queuePosition(state.writerQueue, sessionID); position > 0 {
		return position, nil
	}
	if state.writerID == "" {
		state.writerID = sessionID
		return 0, nil
	}
	state.writerQueue = append(state.writerQueue, sessionID)
	return len(state.writerQueue), nil
}

func (b *Broker) ReleaseWriter(sessionID string) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		return ErrSessionNotFound
	}
	if state.writerID == sessionID {
		state.writerID = ""
		b.promoteWriterLocked(state)
		return nil
	}
	if removeWriterRequest(state, sessionID) {
		return nil
	}
	return ErrWriterRequired
}

func (b *Broker) Write(ctx context.Context, sessionID string, payload []byte) error {
	if len(payload) == 0 || len(payload) > MaxDataFrameBytes {
		return ErrInvalidDataFrame
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	payload = append([]byte(nil), payload...)

	b.mu.Lock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		b.mu.Unlock()
		return ErrSessionNotFound
	}
	if state.writerID != sessionID && !session.request.SharedWrite {
		b.mu.Unlock()
		return ErrWriterRequired
	}
	b.mu.Unlock()

	state.writeMu.Lock()
	defer state.writeMu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	for len(payload) > 0 {
		count, err := state.transport.Write(payload)
		if count > 0 && count <= len(payload) {
			b.metrics.txBytes.Add(uint64(count))
		}
		if err != nil {
			return err
		}
		if count <= 0 || count > len(payload) {
			return io.ErrShortWrite
		}
		payload = payload[count:]
	}
	return nil
}

func (b *Broker) SendBreak(ctx context.Context, sessionID string) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	b.mu.Lock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		b.mu.Unlock()
		return ErrSessionNotFound
	}
	if state.writerID != sessionID && !session.request.SharedWrite {
		b.mu.Unlock()
		return ErrWriterRequired
	}
	transport, ok := state.transport.(breakTransport)
	b.mu.Unlock()
	if !ok {
		return errors.New("serial break is unavailable")
	}
	state.writeMu.Lock()
	defer state.writeMu.Unlock()
	return transport.SendBreak(ctx)
}

func (b *Broker) Detach(sessionID string) error {
	var transport io.Closer
	b.mu.Lock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		b.mu.Unlock()
		return ErrSessionNotFound
	}
	transport = b.detachLocked(state, sessionID, io.EOF)
	b.mu.Unlock()
	if transport != nil {
		return transport.Close()
	}
	return nil
}

func (b *Broker) Session(sessionID string) (SessionSnapshot, bool) {
	b.mu.Lock()
	defer b.mu.Unlock()
	session, state := b.sessionAndChannelLocked(sessionID)
	if session == nil {
		return SessionSnapshot{}, false
	}
	return sessionSnapshot(session, state), true
}

func (b *Broker) Snapshot() BrokerSnapshot {
	b.mu.Lock()
	defer b.mu.Unlock()
	channels := make([]ChannelSnapshot, 0, len(b.channels))
	total := 0
	for _, state := range b.channels {
		queuedBytes := 0
		for _, session := range state.sessions {
			queuedBytes += session.outbound.size()
		}
		total += len(state.sessions)
		channels = append(channels, ChannelSnapshot{
			Channel:          state.channel,
			AttachmentCount:  len(state.sessions),
			MaxAttachments:   b.maxAttachments,
			WriterSessionID:  state.writerID,
			PendingWriterIDs: append([]string(nil), state.writerQueue...),
			QueuedBytes:      queuedBytes,
		})
	}
	sort.Slice(channels, func(left, right int) bool { return channels[left].Channel < channels[right].Channel })
	return BrokerSnapshot{
		TotalAttachments: total, ActiveChannels: len(channels), Channels: channels,
		AttachAttempts:          b.metrics.attachAttempts.Load(),
		SuccessfulAttachments:   b.metrics.successfulAttachments.Load(),
		RejectedAttachments:     b.metrics.rejectedAttachments.Load(),
		TransportOpens:          b.metrics.transportOpens.Load(),
		TransportOpenFailures:   b.metrics.transportOpenFailures.Load(),
		TransportReconnects:     b.metrics.transportReconnects.Load(),
		TransportFailures:       b.metrics.transportFailures.Load(),
		SlowConsumerDisconnects: b.metrics.slowConsumerDisconnects.Load(),
		RXBytes:                 b.metrics.rxBytes.Load(), TXBytes: b.metrics.txBytes.Load(),
		FanoutBytes: b.metrics.fanoutBytes.Load(),
	}
}

func (b *Broker) Close() error {
	b.cancel()
	b.mu.Lock()
	if b.closed {
		b.mu.Unlock()
		return nil
	}
	b.closed = true
	transports := make([]io.Closer, 0, len(b.channels))
	for channel, state := range b.channels {
		for sessionID, session := range state.sessions {
			delete(b.sessions, sessionID)
			session.outbound.close(ErrChannelClosed)
		}
		delete(b.channels, channel)
		transports = append(transports, state.transport)
	}
	b.mu.Unlock()
	var result error
	for _, transport := range transports {
		result = errors.Join(result, transport.Close())
	}
	return result
}

func (b *Broker) readChannel(state *channelState) {
	buffer := make([]byte, serialReadBufferBytes)
	for {
		count, err := state.transport.Read(buffer)
		if count > 0 {
			b.fanOut(state, buffer[:count])
		}
		if err != nil {
			b.failChannel(state, err)
			return
		}
		if count == 0 {
			b.failChannel(state, io.ErrNoProgress)
			return
		}
	}
}

func (b *Broker) fanOut(state *channelState, payload []byte) {
	var transports []io.Closer
	b.mu.Lock()
	if b.channels[state.channel] != state {
		b.mu.Unlock()
		return
	}
	b.metrics.rxBytes.Add(uint64(len(payload)))
	slow := make([]string, 0)
	for sessionID, session := range state.sessions {
		if !session.outbound.push(payload) {
			slow = append(slow, sessionID)
		} else {
			b.metrics.fanoutBytes.Add(uint64(len(payload)))
		}
	}
	b.metrics.slowConsumerDisconnects.Add(uint64(len(slow)))
	for _, sessionID := range slow {
		if transport := b.detachLocked(state, sessionID, ErrSlowConsumer); transport != nil {
			transports = append(transports, transport)
		}
	}
	b.mu.Unlock()
	for _, transport := range transports {
		_ = transport.Close()
	}
}

func (b *Broker) failChannel(state *channelState, cause error) {
	if cause == nil || errors.Is(cause, io.EOF) {
		cause = ErrChannelClosed
	}
	b.mu.Lock()
	if b.channels[state.channel] != state {
		b.mu.Unlock()
		return
	}
	b.metrics.transportFailures.Add(1)
	for sessionID, session := range state.sessions {
		delete(b.sessions, sessionID)
		session.outbound.close(cause)
	}
	delete(b.channels, state.channel)
	b.mu.Unlock()
	_ = state.transport.Close()
}

func (b *Broker) sessionAndChannelLocked(sessionID string) (*brokerSession, *channelState) {
	session := b.sessions[sessionID]
	if session == nil {
		return nil, nil
	}
	return session, b.channels[session.request.Channel]
}

func (b *Broker) detachLocked(state *channelState, sessionID string, cause error) io.Closer {
	session := state.sessions[sessionID]
	if session == nil {
		return nil
	}
	delete(state.sessions, sessionID)
	delete(b.sessions, sessionID)
	removeWriterRequest(state, sessionID)
	if state.writerID == sessionID {
		state.writerID = ""
		b.promoteWriterLocked(state)
	}
	session.outbound.close(cause)
	if len(state.sessions) == 0 {
		delete(b.channels, state.channel)
		return state.transport
	}
	return nil
}

func (b *Broker) promoteWriterLocked(state *channelState) {
	for len(state.writerQueue) > 0 {
		next := state.writerQueue[0]
		state.writerQueue = state.writerQueue[1:]
		if _, exists := state.sessions[next]; exists {
			state.writerID = next
			return
		}
	}
}

func removeWriterRequest(state *channelState, sessionID string) bool {
	for index, queued := range state.writerQueue {
		if queued == sessionID {
			state.writerQueue = append(state.writerQueue[:index], state.writerQueue[index+1:]...)
			return true
		}
	}
	return false
}

func queuePosition(queue []string, sessionID string) int {
	for index, queued := range queue {
		if queued == sessionID {
			return index + 1
		}
	}
	return 0
}

func sessionSnapshot(session *brokerSession, state *channelState) SessionSnapshot {
	return SessionSnapshot{
		SessionID:     session.request.SessionID,
		Channel:       session.request.Channel,
		Actor:         session.request.Actor,
		Writer:        session.request.SharedWrite || state.writerID == session.request.SessionID,
		QueuePosition: queuePosition(state.writerQueue, session.request.SessionID),
		OutboundBytes: session.outbound.size(),
	}
}

func (a *Attachment) SessionID() string {
	return a.sessionID
}

func (a *Attachment) Receive(ctx context.Context) ([]byte, error) {
	return a.outbound.pop(ctx)
}

func (a *Attachment) Write(ctx context.Context, payload []byte) error {
	return a.broker.Write(ctx, a.sessionID, payload)
}

func (a *Attachment) SendBreak(ctx context.Context) error {
	return a.broker.SendBreak(ctx, a.sessionID)
}

func (a *Attachment) RequestWriter() (int, error) {
	return a.broker.RequestWriter(a.sessionID)
}

func (a *Attachment) ReleaseWriter() error {
	return a.broker.ReleaseWriter(a.sessionID)
}

func (a *Attachment) Close() error {
	a.closeOnce.Do(func() {
		a.closeErr = a.broker.Detach(a.sessionID)
		if errors.Is(a.closeErr, ErrSessionNotFound) {
			a.closeErr = nil
		}
	})
	return a.closeErr
}

type byteQueue struct {
	mu       sync.Mutex
	chunks   [][]byte
	bytes    int
	max      int
	notify   chan struct{}
	closed   bool
	closeErr error
}

func newByteQueue(maxBytes int) *byteQueue {
	return &byteQueue{max: maxBytes, notify: make(chan struct{}, 1)}
}

func (q *byteQueue) push(payload []byte) bool {
	q.mu.Lock()
	defer q.mu.Unlock()
	if q.closed || len(payload) == 0 || len(payload) > q.max || q.bytes+len(payload) > q.max {
		return false
	}
	q.chunks = append(q.chunks, append([]byte(nil), payload...))
	q.bytes += len(payload)
	q.signalLocked()
	return true
}

func (q *byteQueue) pop(ctx context.Context) ([]byte, error) {
	for {
		q.mu.Lock()
		if len(q.chunks) > 0 {
			payload := q.chunks[0]
			q.chunks[0] = nil
			q.chunks = q.chunks[1:]
			q.bytes -= len(payload)
			q.mu.Unlock()
			return payload, nil
		}
		if q.closed {
			err := q.closeErr
			if err == nil {
				err = io.EOF
			}
			q.mu.Unlock()
			return nil, err
		}
		notify := q.notify
		q.mu.Unlock()
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-notify:
		}
	}
}

func (q *byteQueue) close(cause error) {
	q.mu.Lock()
	if !q.closed {
		q.closed = true
		q.closeErr = cause
		q.signalLocked()
	}
	q.mu.Unlock()
}

func (q *byteQueue) size() int {
	q.mu.Lock()
	defer q.mu.Unlock()
	return q.bytes
}

func (q *byteQueue) signalLocked() {
	select {
	case q.notify <- struct{}{}:
	default:
	}
}