Upload files to "rtc-c2/pkg/transport"

2026-06-10 17:20:02 +00:00
parent fff1d686c1
commit 4568bfa5c3
2 changed files with 530 additions and 0 deletions
@@ -0,0 +1,104 @@
+// Package transport provides WebRTC-based transport for the rtc-c2 framework.
+// It wraps Pion WebRTC to manage peer connections, data channels, and ICE/TURN connectivity.
+package transport
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"time"
+)
+
+// Role indicates whether this instance creates the offer (Operator) or answers (Beacon).
+type Role int
+
+const (
+	RoleOperator Role = iota
+	RoleBeacon
+)
+
+func (r Role) String() string {
+	switch r {
+	case RoleOperator:
+		return "operator"
+	case RoleBeacon:
+		return "beacon"
+	default:
+		return "unknown"
+	}
+}
+
+// ICEServerConfig defines a STUN or TURN server for NAT traversal.
+type ICEServerConfig struct {
+	URLs       []string `json:"urls"`
+	Username   string   `json:"username,omitempty"`
+	Credential string   `json:"credential,omitempty"`
+}
+
+// Config holds all transport layer configuration.
+type Config struct {
+	// Role of this peer (operator or beacon)
+	Role Role `json:"role"`
+
+	// ICE servers for NAT traversal
+	ICEServers []ICEServerConfig `json:"ice_servers"`
+
+	// Data channel configuration
+	DataChannelName string `json:"data_channel_name"`
+
+	// WebRTC settings
+	Mtu             int           `json:"mtu"`
+	ReconnectDelay  time.Duration `json:"reconnect_delay"`
+	MaxReconnectRetries int       `json:"max_reconnect_retries"`
+
+	// Peer ID — unique identifier for this instance
+	PeerID string `json:"peer_id"`
+}
+
+// DefaultConfig returns a sensible default configuration.
+// Uses Google's public STUN servers by default; TURN can be added.
+func DefaultConfig(role Role) *Config {
+	peerID := generatePeerID()
+
+	return &Config{
+		Role: role,
+		ICEServers: []ICEServerConfig{
+			{
+				URLs: []string{
+					"stun:stun.l.google.com:19302",
+					"stun:stun1.l.google.com:19302",
+				},
+			},
+		},
+		DataChannelName:     "rtc-c2-channel",
+		Mtu:                 1500,
+		ReconnectDelay:      5 * time.Second,
+		MaxReconnectRetries: 10,
+		PeerID:              peerID,
+	}
+}
+
+// WithTURN adds TURN relay servers to the configuration.
+// For development, use coturn. For Ghost Calls-style operation, use provider TURN.
+func (c *Config) WithTURN(urls []string, username, credential string) *Config {
+	c.ICEServers = append(c.ICEServers, ICEServerConfig{
+		URLs:       urls,
+		Username:   username,
+		Credential: credential,
+	})
+	return c
+}
+
+// WithAdditionalSTUN adds STUN servers.
+func (c *Config) WithAdditionalSTUN(urls ...string) *Config {
+	c.ICEServers = append(c.ICEServers, ICEServerConfig{URLs: urls})
+	return c
+}
+
+func generatePeerID() string {
+	b := make([]byte, 8)
+	if _, err := rand.Read(b); err != nil {
+		return fmt.Sprintf("peer-%d", time.Now().UnixNano())
+	}
+	return hex.EncodeToString(b)
+}
@@ -0,0 +1,426 @@
+package transport
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"sync"
+	"time"
+
+	"github.com/pion/webrtc/v4"
+)
+
+// PeerState describes the current connection state.
+type PeerState int
+
+const (
+	StateDisconnected PeerState = iota
+	StateConnecting
+	StateConnected
+	StateFailed
+)
+
+func (s PeerState) String() string {
+	switch s {
+	case StateDisconnected:
+		return "disconnected"
+	case StateConnecting:
+		return "connecting"
+	case StateConnected:
+		return "connected"
+	case StateFailed:
+		return "failed"
+	default:
+		return "unknown"
+	}
+}
+
+// PeerEvent is emitted when the peer connection state changes.
+type PeerEvent struct {
+	State PeerState
+	Err   error
+}
+
+// Peer wraps a Pion RTCPeerConnection and provides a high-level interface
+// for creating and managing WebRTC data channels with automatic reconnection.
+type Peer struct {
+	mu       sync.RWMutex
+	config   *Config
+	pc       *webrtc.PeerConnection
+	dc       *webrtc.DataChannel
+	state    PeerState
+	stopCh   chan struct{}
+	doneCh   chan struct{}
+
+	// Callbacks
+	onStateChange func(PeerEvent)
+	onMessage     func([]byte)
+	onError       func(error)
+
+	// Reconnection
+	reconnectCount int
+
+	// SDP signalling hooks
+	// OnLocalDescription is called when an SDP offer/answer is ready to send.
+	OnLocalDescription func(sdp webrtc.SessionDescription) error
+	// RemoteDescriptionChan receives SDP answers/offers from the remote peer.
+	RemoteDescriptionChan chan webrtc.SessionDescription
+}
+
+// NewPeer creates a new WebRTC peer.
+func NewPeer(cfg *Config) (*Peer, error) {
+	if cfg == nil {
+		cfg = DefaultConfig(RoleBeacon)
+	}
+
+	p := &Peer{
+		config:                cfg,
+		state:                 StateDisconnected,
+		stopCh:                make(chan struct{}),
+		doneCh:                make(chan struct{}),
+		RemoteDescriptionChan: make(chan webrtc.SessionDescription, 16),
+	}
+
+	return p, nil
+}
+
+// OnStateChange registers a callback for connection state changes.
+func (p *Peer) OnStateChange(fn func(PeerEvent)) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.onStateChange = fn
+}
+
+// OnMessage registers a callback for incoming data channel messages.
+func (p *Peer) OnMessage(fn func([]byte)) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.onMessage = fn
+}
+
+// OnError registers a callback for errors.
+func (p *Peer) OnError(fn func(error)) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.onError = fn
+}
+
+// Connect initiates the WebRTC connection.
+// For the operator role, it creates an offer.
+// For the beacon role, it waits for a remote offer.
+func (p *Peer) Connect(ctx context.Context) error {
+	p.mu.Lock()
+	if p.state == StateConnected {
+		p.mu.Unlock()
+		return fmt.Errorf("peer: already connected")
+	}
+	p.state = StateConnecting
+	p.mu.Unlock()
+
+	p.emitState(StateConnecting, nil)
+
+	if err := p.createPeerConnection(); err != nil {
+		p.emitState(StateFailed, err)
+		return fmt.Errorf("peer: create pc: %w", err)
+	}
+
+	// Create data channel (operator initiates)
+	if p.config.Role == RoleOperator {
+		if err := p.createDataChannel(); err != nil {
+			return fmt.Errorf("peer: create dc: %w", err)
+		}
+
+		// Create and send offer
+		offer, err := p.pc.CreateOffer(nil)
+		if err != nil {
+			return fmt.Errorf("peer: create offer: %w", err)
+		}
+
+		if err := p.pc.SetLocalDescription(offer); err != nil {
+			return fmt.Errorf("peer: set local desc: %w", err)
+		}
+
+		// Signal the offer to the remote peer
+		if p.OnLocalDescription != nil {
+			if err := p.OnLocalDescription(offer); err != nil {
+				return fmt.Errorf("peer: signal offer: %w", err)
+			}
+		}
+	}
+
+	go p.connectionLoop(ctx)
+	return nil
+}
+
+// Send sends data over the data channel.
+func (p *Peer) Send(data []byte) error {
+	p.mu.RLock()
+	dc := p.dc
+	p.mu.RUnlock()
+
+	if dc == nil {
+		return fmt.Errorf("peer: data channel not ready")
+	}
+
+	return dc.Send(data)
+}
+
+// Close terminates the peer connection.
+func (p *Peer) Close() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	select {
+	case <-p.stopCh:
+		return nil
+	default:
+		close(p.stopCh)
+	}
+
+	if p.pc != nil {
+		return p.pc.Close()
+	}
+	return nil
+}
+
+// Done returns a channel that closes when the peer is fully shut down.
+func (p *Peer) Done() <-chan struct{} {
+	return p.doneCh
+}
+
+// Config returns the peer's configuration.
+func (p *Peer) Config() *Config {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.config
+}
+
+// State returns the current peer state.
+func (p *Peer) State() PeerState {
+	p.mu.RLock()
+	defer p.mu.RUnlock()
+	return p.state
+}
+
+// createPeerConnection sets up the underlying RTCPeerConnection.
+func (p *Peer) createPeerConnection() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	settings := webrtc.SettingEngine{}
+	settings.DetachDataChannels()
+
+	api := webrtc.NewAPI(webrtc.WithSettingEngine(settings))
+
+	iceServers := make([]webrtc.ICEServer, len(p.config.ICEServers))
+	for i, s := range p.config.ICEServers {
+		iceServers[i] = webrtc.ICEServer{
+			URLs:       s.URLs,
+			Username:   s.Username,
+			Credential: s.Credential,
+		}
+	}
+
+	config := webrtc.Configuration{
+		ICEServers:   iceServers,
+		ICETransportPolicy: webrtc.ICETransportPolicyAll,
+	}
+
+	pc, err := api.NewPeerConnection(config)
+	if err != nil {
+		return fmt.Errorf("new pc: %w", err)
+	}
+
+	// Handle ICE connection state changes
+	pc.OnICEConnectionStateChange(func(state webrtc.ICEConnectionState) {
+		log.Printf("[transport] ICE state: %s", state)
+		switch state {
+		case webrtc.ICEConnectionStateConnected:
+			p.emitState(StateConnected, nil)
+			p.mu.Lock()
+			p.reconnectCount = 0
+			p.mu.Unlock()
+		case webrtc.ICEConnectionStateDisconnected:
+			p.emitState(StateDisconnected, nil)
+		case webrtc.ICEConnectionStateFailed:
+			p.emitState(StateFailed, fmt.Errorf("ICE failed"))
+		}
+	})
+
+	// Handle data channels initiated by the remote peer (beacon side)
+	if p.config.Role == RoleBeacon {
+		pc.OnDataChannel(func(dc *webrtc.DataChannel) {
+			log.Printf("[transport] received remote data channel: %s", dc.Label())
+			p.handleDataChannel(dc)
+		})
+	}
+
+	p.pc = pc
+	return nil
+}
+
+// createDataChannel creates the data channel (operator side).
+func (p *Peer) createDataChannel() error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	dc, err := p.pc.CreateDataChannel(p.config.DataChannelName, &webrtc.DataChannelInit{
+		Ordered: boolPtr(true),
+	})
+	if err != nil {
+		return fmt.Errorf("create dc: %w", err)
+	}
+
+	p.handleDataChannel(dc)
+	return nil
+}
+
+// handleDataChannel sets up the data channel callbacks.
+func (p *Peer) handleDataChannel(dc *webrtc.DataChannel) {
+	p.mu.Lock()
+	p.dc = dc
+	p.mu.Unlock()
+
+	dc.OnOpen(func() {
+		log.Printf("[transport] data channel '%s' opened", dc.Label())
+	})
+
+	dc.OnClose(func() {
+		log.Printf("[transport] data channel '%s' closed", dc.Label())
+	})
+
+	dc.OnMessage(func(msg webrtc.DataChannelMessage) {
+		p.mu.RLock()
+		cb := p.onMessage
+		p.mu.RUnlock()
+		if cb != nil {
+			cb(msg.Data)
+		}
+	})
+}
+
+// connectionLoop handles the main event loop including reconnection.
+func (p *Peer) connectionLoop(ctx context.Context) {
+	defer close(p.doneCh)
+
+	for {
+		select {
+		case <-ctx.Done():
+			p.Close()
+			return
+
+		case <-p.stopCh:
+			return
+
+		case remoteDesc := <-p.RemoteDescriptionChan:
+			if err := p.pc.SetRemoteDescription(remoteDesc); err != nil {
+				log.Printf("[transport] set remote desc error: %v", err)
+				p.emitState(StateFailed, err)
+				return
+			}
+
+			// If we're the beacon, we need to create an answer
+			if p.config.Role == RoleBeacon {
+				answer, err := p.pc.CreateAnswer(nil)
+				if err != nil {
+					log.Printf("[transport] create answer error: %v", err)
+					return
+				}
+
+				if err := p.pc.SetLocalDescription(answer); err != nil {
+					log.Printf("[transport] set local desc error: %v", err)
+					return
+				}
+
+				if p.OnLocalDescription != nil {
+					if err := p.OnLocalDescription(answer); err != nil {
+						log.Printf("[transport] signal answer error: %v", err)
+					}
+				}
+			}
+
+		case <-p.reconnectTimer():
+			p.mu.RLock()
+			count := p.reconnectCount
+			maxRetries := p.config.MaxReconnectRetries
+			p.mu.RUnlock()
+
+			if count >= maxRetries {
+				p.emitState(StateFailed, fmt.Errorf("max reconnection retries exceeded"))
+				return
+			}
+
+			if p.state == StateDisconnected || p.state == StateFailed {
+				log.Printf("[transport] attempting reconnect #%d", count+1)
+				p.mu.Lock()
+				p.reconnectCount++
+				p.mu.Unlock()
+				p.reconnect()
+			}
+		}
+	}
+}
+
+func (p *Peer) reconnectTimer() <-chan time.Time {
+	p.mu.RLock()
+	delay := p.config.ReconnectDelay
+	p.mu.RUnlock()
+	return time.After(delay)
+}
+
+func (p *Peer) reconnect() {
+	_ = p.pc.Close()
+	p.mu.Lock()
+	p.pc = nil
+	p.dc = nil
+	p.mu.Unlock()
+
+	p.emitState(StateConnecting, nil)
+
+	if err := p.createPeerConnection(); err != nil {
+		log.Printf("[transport] reconnect create pc error: %v", err)
+		p.emitState(StateFailed, err)
+		return
+	}
+
+	if p.config.Role == RoleOperator {
+		if err := p.createDataChannel(); err != nil {
+			log.Printf("[transport] reconnect create dc error: %v", err)
+			return
+		}
+
+		offer, err := p.pc.CreateOffer(nil)
+		if err != nil {
+			log.Printf("[transport] reconnect create offer error: %v", err)
+			return
+		}
+
+		if err := p.pc.SetLocalDescription(offer); err != nil {
+			log.Printf("[transport] reconnect set local error: %v", err)
+			return
+		}
+
+		if p.OnLocalDescription != nil {
+			_ = p.OnLocalDescription(offer)
+		}
+	}
+}
+
+func (p *Peer) emitState(state PeerState, err error) {
+	p.mu.Lock()
+	p.state = state
+	cb := p.onStateChange
+	p.mu.Unlock()
+
+	if cb != nil {
+		cb(PeerEvent{State: state, Err: err})
+	}
+}
+
+// Ensure io.Closer interface
+var _ io.Closer = (*Peer)(nil)
+
+func boolPtr(b bool) *bool {
+	return &b
+}