project structure refactoring

http/rtc/listener.go

@@ -0,0 +1,77 @@
+package rtc
+
+import (
+	"net"
+	"rkkvm/config"
+
+	"github.com/pion/webrtc/v4"
+	log "github.com/sirupsen/logrus"
+)
+
+func initUDPListener(host string, port int) (*net.UDPConn, error) {
+	l, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP(host), Port: port})
+	if err != nil {
+		return nil, ErrWebRTCParam("failed to init webrtc listener: %v", err)
+	}
+
+	// Increase the UDP receive buffer size
+	// Default UDP buffer sizes vary on different operating systems
+	bufferSize := 300000 // 300KB
+	err = l.SetReadBuffer(bufferSize)
+	if err != nil {
+		return nil, ErrWebRTCParam("failed to set read buffer: %v", err)
+	}
+
+	return l, nil
+}
+
+func InitListener(host string, port int, aPort int) (*RTC, error) {
+	vl, err := initUDPListener(host, port)
+	if err != nil {
+		return nil, err
+	}
+
+	al, err := initUDPListener(host, aPort)
+	if err != nil {
+		return nil, err
+	}
+
+	mimeType := ""
+	switch config.Get().Video.Codec {
+	case config.StreamSourceH264:
+		mimeType = webrtc.MimeTypeH264
+	case config.StreamSourceHevc: // WebRTC currently has no official support for H265
+		mimeType = webrtc.MimeTypeH265
+	default:
+		return nil, ErrWebRTCParam("unknown video codec: %s", config.Get().Video.Codec)
+	}
+
+	video, _ := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: mimeType}, "video", "rkkvm")
+	audio, _ := webrtc.NewTrackLocalStaticRTP(webrtc.RTPCodecCapability{MimeType: webrtc.MimeTypeOpus}, "audio", "rkkvm")
+	rtc = &RTC{
+		videoListener: vl,
+		audioListener: al,
+		peers:         make(map[string]*webrtc.PeerConnection),
+		videoTrack:    video,
+		audioTrack:    audio,
+	}
+
+	return rtc, nil
+}
+
+func listenerRead(l *net.UDPConn, track *webrtc.TrackLocalStaticRTP) {
+	buf := make([]byte, 1600) // Buffer to hold incoming RTP packets
+
+	for {
+		n, _, err := l.ReadFrom(buf)
+		if err != nil {
+			log.Errorf("error reading from UDP: %v\n", err)
+			continue
+		}
+
+		_, err = track.Write(buf[:n])
+		if err != nil {
+			log.Errorf("failed to send RTP to peer: %v", err)
+		}
+	}
+}

http/rtc/webrtc.go

@@ -0,0 +1,139 @@
+package rtc
+
+import (
+	"errors"
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/google/uuid"
+	log "github.com/sirupsen/logrus"
+
+	"github.com/pion/webrtc/v4"
+)
+
+var rtc *RTC
+
+func Get() *RTC {
+	return rtc
+}
+
+var ErrWebRTC = errors.New("webrtc")
+var ErrWebRTCParam = func(format string, args ...any) error {
+	return fmt.Errorf("%w: "+format, args...)
+}
+var ErrPeerClosedConn = ErrWebRTCParam("peer closed conn")
+
+type RTC struct {
+	peers         map[string]*webrtc.PeerConnection
+	videoListener *net.UDPConn
+	audioListener *net.UDPConn
+	videoTrack    *webrtc.TrackLocalStaticRTP
+	audioTrack    *webrtc.TrackLocalStaticRTP
+	m             sync.Mutex
+}
+
+func (r *RTC) AddPeer(p *webrtc.PeerConnection, offer webrtc.SessionDescription) (*webrtc.SessionDescription, error) {
+	peerID := uuid.New().String()
+	r.m.Lock()
+	r.peers[peerID] = p
+	r.m.Unlock()
+
+	p.OnConnectionStateChange(func(connState webrtc.PeerConnectionState) {
+		if connState == webrtc.PeerConnectionStateFailed || connState == webrtc.PeerConnectionStateClosed {
+			r.m.Lock()
+			defer r.m.Unlock()
+			delete(r.peers, peerID)
+			p.Close()
+
+			peers := make([]string, 0, len(r.peers))
+			for p := range r.peers {
+				peers = append(peers, p)
+			}
+			log.WithField("peers", peers).Infof("Peer %s disconnected and resources cleaned up.", peerID)
+		}
+	})
+
+	vSender, err := p.AddTrack(r.videoTrack)
+	if err != nil {
+		return nil, ErrWebRTCParam("failed to add video track: %v", err)
+	}
+	processRTCP(vSender)
+	aSender, err := p.AddTrack(r.audioTrack)
+	if err != nil {
+		return nil, ErrWebRTCParam("failed to add audio track: %v", err)
+	}
+	processRTCP(aSender)
+
+	if err := p.SetRemoteDescription(offer); err != nil {
+		return nil, ErrWebRTCParam("failed to set remote description: %v", err)
+	}
+
+	answer, err := p.CreateAnswer(nil)
+	if err != nil {
+		return nil, ErrWebRTCParam("failed to create answer: %v", err)
+	}
+	gatherComplete := webrtc.GatheringCompletePromise(p)
+
+	if err := p.SetLocalDescription(answer); err != nil {
+		return nil, ErrWebRTCParam("failed to set local description: %v", err)
+	}
+	<-gatherComplete
+
+	return p.LocalDescription(), nil
+}
+
+func (r *RTC) VideoListenerRead() {
+	listenerRead(r.videoListener, r.videoTrack)
+}
+
+func (r *RTC) AudioListenerRead() {
+	listenerRead(r.audioListener, r.audioTrack)
+}
+
+func (r *RTC) Close() error {
+	r.videoListener.Close()
+	r.audioListener.Close()
+
+	return nil
+}
+
+func NewPeer() (*webrtc.PeerConnection, error) {
+	peer, err := webrtc.NewPeerConnection(webrtc.Configuration{
+		ICEServers: []webrtc.ICEServer{
+			{
+				URLs: []string{"stun:stun.l.google.com:19302"},
+			},
+		},
+	})
+	if err == nil {
+		// Set the handler for ICE connection state
+		// This will notify you when the peer has connected/disconnected
+		peer.OnICEConnectionStateChange(func(connState webrtc.ICEConnectionState) {
+			log.Infof("Connection State has changed %s", connState.String())
+
+			if connState == webrtc.ICEConnectionStateFailed {
+				if closeErr := peer.Close(); closeErr != nil {
+					panic(closeErr)
+				}
+			}
+		})
+	}
+
+	return peer, err
+}
+
+// Read incoming RTCP packets
+// Before these packets are retuned they are processed by interceptors. For things
+// like NACK this needs to be called.
+func processRTCP(rtpSender *webrtc.RTPSender) {
+	go func() {
+		rtcpBuf := make([]byte, 1500)
+
+		for {
+			if _, _, rtcpErr := rtpSender.Read(rtcpBuf); rtcpErr != nil {
+				return
+			}
+		}
+	}()
+}