From 7cd469f7bdbee84bb09c5c407defd3ccc4345b77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miroslav=20=C5=A0ediv=C3=BD?= <sedivy.miro@gmail.com>
Date: Mon, 10 Apr 2023 21:21:11 +0200
Subject: [PATCH] cleanup peer.

---
 internal/webrtc/manager.go | 17 ++++++++----
 internal/webrtc/peer.go    | 57 ++++++++------------------------------
 2 files changed, 23 insertions(+), 51 deletions(-)

diff --git a/internal/webrtc/manager.go b/internal/webrtc/manager.go
index 29fc8654..f0fd27cb 100644
--- a/internal/webrtc/manager.go
+++ b/internal/webrtc/manager.go
@@ -448,18 +448,23 @@ func (manager *WebRTCManagerCtx) CreatePeer(session types.Session, bitrate int,
 	}
 
 	peer := &WebRTCPeerCtx{
-		logger:                 logger,
-		connection:             connection,
-		dataChannel:            dataChannel,
+		logger:     logger,
+		connection: connection,
+		// tracks & channels
+		audioTrack:  audioTrack,
+		videoTrack:  videoTrack,
+		dataChannel: dataChannel,
+		rtcpChannel: videoRtcp,
+		// config
+		iceTrickle: manager.config.ICETrickle,
+		// deprecated functions
 		changeVideoFromBitrate: changeVideoFromBitrate,
 		changeVideoFromID:      changeVideoFromID,
-		// TODO: Refactor.
-		videoId: videoTrack.stream.ID,
+		videoId:                videoTrack.stream.ID,
 		setPaused: func(isPaused bool) {
 			videoTrack.SetPaused(isPaused)
 			audioTrack.SetPaused(isPaused)
 		},
-		iceTrickle: manager.config.ICETrickle,
 		setVideoAuto: func(videoAuto bool) {
 			// if estimator is enabled and not in passive mode, enable video auto bitrate
 			if manager.config.EstimatorEnabled && !manager.config.EstimatorPassive {
diff --git a/internal/webrtc/peer.go b/internal/webrtc/peer.go
index 09e982fa..7edd3e8f 100644
--- a/internal/webrtc/peer.go
+++ b/internal/webrtc/peer.go
@@ -5,6 +5,7 @@ import (
 	"encoding/binary"
 	"sync"
 
+	"github.com/pion/rtcp"
 	"github.com/pion/webrtc/v3"
 	"github.com/rs/zerolog"
 
@@ -13,27 +14,29 @@ import (
 )
 
 type WebRTCPeerCtx struct {
-	mu                     sync.Mutex
-	logger                 zerolog.Logger
-	connection             *webrtc.PeerConnection
-	dataChannel            *webrtc.DataChannel
+	mu         sync.Mutex
+	logger     zerolog.Logger
+	connection *webrtc.PeerConnection
+	// tracks & channels
+	audioTrack  *Track
+	videoTrack  *Track
+	dataChannel *webrtc.DataChannel
+	rtcpChannel chan []rtcp.Packet
+	// config
+	iceTrickle bool
+	// deprecated functions
 	changeVideoFromBitrate func(bitrate int)
 	changeVideoFromID      func(id string) int
 	videoId                func() string
 	setPaused              func(isPaused bool)
 	setVideoAuto           func(auto bool)
 	getVideoAuto           func() bool
-	iceTrickle             bool
 }
 
 func (peer *WebRTCPeerCtx) CreateOffer(ICERestart bool) (*webrtc.SessionDescription, error) {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return nil, types.ErrWebRTCConnectionNotFound
-	}
-
 	offer, err := peer.connection.CreateOffer(&webrtc.OfferOptions{
 		ICERestart: ICERestart,
 	})
@@ -48,10 +51,6 @@ func (peer *WebRTCPeerCtx) CreateAnswer() (*webrtc.SessionDescription, error) {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return nil, types.ErrWebRTCConnectionNotFound
-	}
-
 	answer, err := peer.connection.CreateAnswer(nil)
 	if err != nil {
 		return nil, err
@@ -83,10 +82,6 @@ func (peer *WebRTCPeerCtx) SetOffer(sdp string) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	return peer.connection.SetRemoteDescription(webrtc.SessionDescription{
 		SDP:  sdp,
 		Type: webrtc.SDPTypeOffer,
@@ -97,10 +92,6 @@ func (peer *WebRTCPeerCtx) SetAnswer(sdp string) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	return peer.connection.SetRemoteDescription(webrtc.SessionDescription{
 		SDP:  sdp,
 		Type: webrtc.SDPTypeAnswer,
@@ -111,10 +102,6 @@ func (peer *WebRTCPeerCtx) SetCandidate(candidate webrtc.ICECandidateInit) error
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	return peer.connection.AddICECandidate(candidate)
 }
 
@@ -122,10 +109,6 @@ func (peer *WebRTCPeerCtx) SetVideoBitrate(peerBitrate int) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	peer.changeVideoFromBitrate(peerBitrate)
 	return nil
 }
@@ -134,10 +117,6 @@ func (peer *WebRTCPeerCtx) SetVideoID(videoID string) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	peer.changeVideoFromID(videoID)
 	return nil
 }
@@ -154,10 +133,6 @@ func (peer *WebRTCPeerCtx) SetPaused(isPaused bool) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.connection == nil {
-		return types.ErrWebRTCConnectionNotFound
-	}
-
 	peer.logger.Info().Bool("is_paused", isPaused).Msg("set paused")
 	peer.setPaused(isPaused)
 	return nil
@@ -167,10 +142,6 @@ func (peer *WebRTCPeerCtx) SendCursorPosition(x, y int) error {
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.dataChannel == nil {
-		return types.ErrWebRTCDataChannelNotFound
-	}
-
 	data := payload.CursorPosition{
 		Header: payload.Header{
 			Event:  payload.OP_CURSOR_POSITION,
@@ -192,10 +163,6 @@ func (peer *WebRTCPeerCtx) SendCursorImage(cur *types.CursorImage, img []byte) e
 	peer.mu.Lock()
 	defer peer.mu.Unlock()
 
-	if peer.dataChannel == nil {
-		return types.ErrWebRTCDataChannelNotFound
-	}
-
 	data := payload.CursorImage{
 		Header: payload.Header{
 			Event:  payload.OP_CURSOR_IMAGE,