server: support UDP

serverudpl.go

@@ -0,0 +1,145 @@
+package gortsplib
+
+import (
+	"net"
+	"sync"
+	"time"
+
+	"github.com/aler9/gortsplib/pkg/multibuffer"
+)
+
+const (
+	// use the same buffer size as gstreamer's rtspsrc
+	kernelReadBufferSize = 0x80000
+
+	readBufferSize = 2048
+)
+
+type publisherData struct {
+	publisher *ServerConn
+	trackID   int
+}
+
+type publisherAddr struct {
+	ip   [net.IPv6len]byte // use a fixed-size array to enable the equality operator
+	port int
+}
+
+func (p *publisherAddr) fill(ip net.IP, port int) {
+	p.port = port
+
+	if len(ip) == net.IPv4len {
+		copy(p.ip[0:], []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff, 0xff}) // v4InV6Prefix
+		copy(p.ip[12:], ip)
+	} else {
+		copy(p.ip[:], ip)
+	}
+}
+
+// ServerUDPListener is a UDP server that can be used to send and receive RTP and RTCP packets.
+type ServerUDPListener struct {
+	streamType StreamType
+
+	pc              *net.UDPConn
+	readBuf         *multibuffer.MultiBuffer
+	publishersMutex sync.RWMutex
+	publishers      map[publisherAddr]*publisherData
+	writeMutex      sync.Mutex
+
+	// out
+	done chan struct{}
+}
+
+// NewServerUDPListener allocates a ServerUDPListener.
+func NewServerUDPListener(address string) (*ServerUDPListener, error) {
+	tmp, err := net.ListenPacket("udp", address)
+	if err != nil {
+		return nil, err
+	}
+	pc := tmp.(*net.UDPConn)
+
+	err = pc.SetReadBuffer(kernelReadBufferSize)
+	if err != nil {
+		return nil, err
+	}
+
+	s := &ServerUDPListener{
+		pc:         pc,
+		readBuf:    multibuffer.New(1, readBufferSize),
+		publishers: make(map[publisherAddr]*publisherData),
+		done:       make(chan struct{}),
+	}
+
+	go s.run()
+
+	return s, nil
+}
+
+// Close closes the listener.
+func (s *ServerUDPListener) Close() {
+	s.pc.Close()
+	<-s.done
+}
+
+func (s *ServerUDPListener) run() {
+	defer close(s.done)
+
+	for {
+		buf := s.readBuf.Next()
+		n, addr, err := s.pc.ReadFromUDP(buf)
+		if err != nil {
+			break
+		}
+
+		func() {
+			s.publishersMutex.RLock()
+			defer s.publishersMutex.RUnlock()
+
+			// find publisher data
+			var pubAddr publisherAddr
+			pubAddr.fill(addr.IP, addr.Port)
+			pubData, ok := s.publishers[pubAddr]
+			if !ok {
+				return
+			}
+
+			pubData.publisher.readHandlers.OnFrame(pubData.trackID, s.streamType, buf[:n])
+		}()
+	}
+}
+
+func (s *ServerUDPListener) port() int {
+	return s.pc.LocalAddr().(*net.UDPAddr).Port
+}
+
+func (s *ServerUDPListener) write(writeTimeout time.Duration, buf []byte, addr *net.UDPAddr) error {
+	s.writeMutex.Lock()
+	defer s.writeMutex.Unlock()
+
+	s.pc.SetWriteDeadline(time.Now().Add(writeTimeout))
+	_, err := s.pc.WriteTo(buf, addr)
+	return err
+}
+
+func (s *ServerUDPListener) addPublisher(ip net.IP, port int, trackID int, sc *ServerConn) {
+	s.publishersMutex.Lock()
+	defer s.publishersMutex.Unlock()
+
+	var addr publisherAddr
+	addr.fill(ip, port)
+
+	s.publishers[addr] = &publisherData{
+		publisher: sc,
+		trackID:   trackID,
+	}
+}
+
+func (s *ServerUDPListener) removePublisher(ip net.IP, port int) {
+	s.publishersMutex.Lock()
+	defer s.publishersMutex.Unlock()
+
+	var addr publisherAddr
+	addr.fill(ip, port)
+
+	delete(s.publishers, addr)
+}