133 lines
2.9 KiB
Go
133 lines
2.9 KiB
Go
package webrtc
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"sync"
|
|
|
|
"github.com/pion/rtp"
|
|
)
|
|
|
|
// Source reads RTP packets from a local UDP socket and fans them out to
|
|
// subscribed peers via per-subscriber buffered channels.
|
|
type Source struct {
|
|
id string
|
|
conn *net.UDPConn
|
|
|
|
mu sync.Mutex
|
|
subscribers map[chan *rtp.Packet]struct{}
|
|
started bool
|
|
closed bool
|
|
done chan struct{}
|
|
}
|
|
|
|
// NewSource binds a UDP socket on 127.0.0.1:port. Pass port=0 to let the OS
|
|
// assign an ephemeral port (useful for tests).
|
|
func NewSource(streamID string, port int) (*Source, error) {
|
|
addr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: port}
|
|
conn, err := net.ListenUDP("udp4", addr)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("webrtc: listen udp: %w", err)
|
|
}
|
|
return &Source{
|
|
id: streamID,
|
|
conn: conn,
|
|
subscribers: make(map[chan *rtp.Packet]struct{}),
|
|
done: make(chan struct{}),
|
|
}, nil
|
|
}
|
|
|
|
// ID returns the registered stream identifier.
|
|
func (s *Source) ID() string { return s.id }
|
|
|
|
// LocalAddr returns the UDP address the source is listening on.
|
|
func (s *Source) LocalAddr() *net.UDPAddr {
|
|
return s.conn.LocalAddr().(*net.UDPAddr)
|
|
}
|
|
|
|
// Subscribe returns a new buffered channel that receives every RTP packet
|
|
// read from the UDP socket. bufDepth is the channel buffer size; when full,
|
|
// packets are dropped (preventing a slow subscriber from back-pressuring
|
|
// the reader).
|
|
func (s *Source) Subscribe(bufDepth int) chan *rtp.Packet {
|
|
ch := make(chan *rtp.Packet, bufDepth)
|
|
s.mu.Lock()
|
|
s.subscribers[ch] = struct{}{}
|
|
s.mu.Unlock()
|
|
return ch
|
|
}
|
|
|
|
// Unsubscribe removes ch from the subscriber set and closes it.
|
|
func (s *Source) Unsubscribe(ch chan *rtp.Packet) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if _, ok := s.subscribers[ch]; ok {
|
|
delete(s.subscribers, ch)
|
|
close(ch)
|
|
}
|
|
}
|
|
|
|
// Start begins the RTP reader goroutine. Safe to call once; subsequent calls
|
|
// are no-ops.
|
|
func (s *Source) Start() {
|
|
s.mu.Lock()
|
|
if s.started || s.closed {
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
s.started = true
|
|
s.mu.Unlock()
|
|
|
|
go s.readLoop()
|
|
}
|
|
|
|
func (s *Source) readLoop() {
|
|
buf := make([]byte, 1500) // MTU-sized; RTP over UDP should fit
|
|
for {
|
|
select {
|
|
case <-s.done:
|
|
return
|
|
default:
|
|
}
|
|
|
|
n, _, err := s.conn.ReadFromUDP(buf)
|
|
if err != nil {
|
|
// Socket closed or error — exit the loop.
|
|
return
|
|
}
|
|
|
|
pkt := &rtp.Packet{}
|
|
if err := pkt.Unmarshal(buf[:n]); err != nil {
|
|
// Malformed packet; skip without crashing.
|
|
continue
|
|
}
|
|
|
|
s.mu.Lock()
|
|
for ch := range s.subscribers {
|
|
select {
|
|
case ch <- pkt:
|
|
default:
|
|
// Subscriber full — drop to protect the reader.
|
|
}
|
|
}
|
|
s.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
// Close stops the reader goroutine, closes the UDP socket, and closes every
|
|
// subscriber channel.
|
|
func (s *Source) Close() error {
|
|
s.mu.Lock()
|
|
if s.closed {
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
s.closed = true
|
|
close(s.done)
|
|
for ch := range s.subscribers {
|
|
delete(s.subscribers, ch)
|
|
close(ch)
|
|
}
|
|
s.mu.Unlock()
|
|
return s.conn.Close()
|
|
}
|