TeleportFling/internal/output/sender.go

// Package output implements the TCP sender and receiver-connection manager
// for the Teleport protocol. It mirrors the design of obs-teleport's
// Sender: a per-connection buffered channel (capacity 1000), frame drop
// when the queue exceeds 800, and a warning at 100.
//
// Usage:
//
//	s := output.New()
//	s.Listen(":9756")
//	s.Send(packet)           // from any goroutine
//	...
//	s.Close()
package output

import (
	"log"
	"net"
	"sync"
	"sync/atomic"
)

const (
	sendChanCap = 1000
	dropAt      = 800
	warnAt      = 100
)

// Sender manages a TCP listener and a set of connected receivers.
// Frames are distributed to every connected receiver; a receiver whose
// queue grows too large has its next frame silently dropped.
type Sender struct {
	mu    sync.Mutex
	conns map[net.Conn]chan []byte
	wg    sync.WaitGroup

	listener net.Listener
	port     int

	dropped atomic.Int64
	bytes   atomic.Int64
}

// New creates an unconnected Sender.
func New() *Sender {
	return &Sender{
		conns: make(map[net.Conn]chan []byte),
	}
}

// Listen binds a TCP listener on addr (e.g. ":9756") and starts accepting
// connections. It returns the resolved port number on success. The caller
// must eventually call Close.
func (s *Sender) Listen(addr string) (int, error) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return 0, err
	}
	s.listener = l

	// Extract the port so the caller can advertise it via UDP announce.
	_, p, err := net.SplitHostPort(l.Addr().String())
	if err != nil {
		if cerr := l.Close(); cerr != nil {
			log.Printf("output: close failed: %v", cerr)
		}
		return 0, err
	}

	var port int
	for _, ch := range p {
		port = port*10 + int(ch-'0')
	}
	s.port = port

	s.wg.Add(1)
	go s.acceptLoop()

	return port, nil
}

// Port returns the TCP port that was assigned at Listen time.
func (s *Sender) Port() int { return s.port }

// NumConns returns the number of connected receivers.
func (s *Sender) NumConns() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.conns)
}

// Send broadcasts a serialised packet to all connected receivers. If a
// receiver's buffered channel is full (> dropAt) the frame is silently
// dropped; a warning is logged at warnAt. Dropped frames are counted and
// exposed via Dropped.
func (s *Sender) Send(b []byte) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Count the packet once as produced bandwidth (independent of how many
	// receivers are attached).
	s.bytes.Add(int64(len(b)))

	for c, ch := range s.conns {
		switch {
		case len(ch) > dropAt:
			log.Printf("output: drop [%s] (queue %d)", c.RemoteAddr(), len(ch))
		case len(ch) > warnAt:
			log.Printf("output: high [%s] (queue %d)", c.RemoteAddr(), len(ch))
		default:
		}

		if len(ch) > dropAt {
			s.dropped.Add(1)
			continue
		}

		// Non-blocking send: if the channel is full after the drop
		// window, drop the frame.
		select {
		case ch <- b:
		default:
			s.dropped.Add(1)
			log.Printf("output: drop [%s] (queue full)", c.RemoteAddr())
		}
	}
}

// Dropped returns the total number of frames dropped due to backpressure.
func (s *Sender) Dropped() int64 {
	return s.dropped.Load()
}

// BytesSent returns the total number of payload bytes handed to Send.
func (s *Sender) BytesSent() int64 {
	return s.bytes.Load()
}

// Close shuts down the listener and waits for all writer goroutines to
// drain. After Close returns the Sender must not be reused.
func (s *Sender) Close() {
	if s.listener != nil {
		// Best-effort close; the accept loop will observe the listener error.
		_ = s.listener.Close()
	}

	s.mu.Lock()
	for _, ch := range s.conns {
		close(ch)
	}
	s.conns = nil
	s.mu.Unlock()

	s.wg.Wait()
}

// acceptLoop runs in its own goroutine and adds incoming connections.
func (s *Sender) acceptLoop() {
	defer s.wg.Done()
	for {
		c, err := s.listener.Accept()
		if err != nil {
			return
		}
		s.SenderAdd(c)
	}
}

// SenderAdd registers a new receiver connection and spawns its writer.
func (s *Sender) SenderAdd(c net.Conn) {
	s.mu.Lock()
	defer s.mu.Unlock()

	log.Printf("output: connect %s", c.RemoteAddr())

	ch := make(chan []byte, sendChanCap)
	s.conns[c] = ch

	s.wg.Add(1)
	go func() {
		defer s.wg.Done()
		defer func() {
			// Best-effort close; writer loop also closes on write error.
			_ = c.Close()
		}()

		for b := range ch {
			if _, err := c.Write(b); err != nil {
				log.Printf("output: disconnect %s", c.RemoteAddr())
				s.removeConn(c)
				return
			}
		}
	}()
}

func (s *Sender) removeConn(c net.Conn) {
	s.mu.Lock()
	delete(s.conns, c)
	s.mu.Unlock()
}