From 8a368bfea3bb4a4979cc67318fc000620a5766d4 Mon Sep 17 00:00:00 2001
From: boks1971 <raja.gobi@tutanota.com>
Date: Thu, 31 Oct 2024 00:26:27 +0530
Subject: [PATCH] Add option to disble peer reflexive discovery

There could be a mismatch between the two ends in candidate priority
when using peer reflexive. It happens in the following scenario

1. Client has two srflx candidates.
   a. The first one gets discovered by LiveKit server as prflx.
   b. The second one gets added via ice-trickle first and then
      gets a STUN ping. So, it is srflx remote candidate from
      server's point-of-view.
2. This leads to a priority issue.
   a. Both candidates have same priority from client's point-of-view
      (both are srflx).
   b. But, from server's point-of-view, the first candidate has
      higher priority (prflx).
3. The first candidate establishes connectivity and becomes
   the selected pair (client is ICE controlling and server is
   ICE controlled, server is in ICE lite).
4. libwebrtc does a sort and switch some time later based on RTT.
   As client side has both at same priority, RTT based sorting
   could make the second candidate the preferred one.
   So, the client sends useCandidate=1 for the second candidate.
   pion/ice does not switch because the selected pair is at
   higher priority due to prflx candidate.
5. STUN pings do not happen and the ICE connection eventually fails.
---
 agent.go        | 10 ++++++++--
 agent_config.go |  4 ++++
 agent_test.go   | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/agent.go b/agent.go
index b9268b8b..27760093 100644
--- a/agent.go
+++ b/agent.go
@@ -145,6 +145,8 @@ type Agent struct {
 	insecureSkipVerify bool
 
 	proxyDialer proxy.Dialer
+
+	disablePeerReflexiveDiscovery bool
 }
 
 // NewAgent creates a new Agent
@@ -219,6 +221,8 @@ func NewAgent(config *AgentConfig) (*Agent, error) { //nolint:gocognit
 		disableActiveTCP: config.DisableActiveTCP,
 
 		userBindingRequestHandler: config.BindingRequestHandler,
+
+		disablePeerReflexiveDiscovery: config.DisablePeerReflexiveDiscovery,
 	}
 	a.connectionStateNotifier = &handlerNotifier{connectionStateFunc: a.onConnectionStateChange, done: make(chan struct{})}
 	a.candidateNotifier = &handlerNotifier{candidateFunc: a.onCandidate, done: make(chan struct{})}
@@ -1039,7 +1043,7 @@ func (a *Agent) handleInbound(m *stun.Message, local Candidate, remote net.Addr)
 			return
 		}
 
-		if remoteCandidate == nil {
+		if remoteCandidate == nil && !a.disablePeerReflexiveDiscovery {
 			ip, port, networkType, err := parseAddr(remote)
 			if err != nil {
 				a.log.Errorf("Failed to create parse remote net.Addr when creating remote prflx candidate: %s", err)
@@ -1066,7 +1070,9 @@ func (a *Agent) handleInbound(m *stun.Message, local Candidate, remote net.Addr)
 			a.addRemoteCandidate(remoteCandidate)
 		}
 
-		a.selector.HandleBindingRequest(m, local, remoteCandidate)
+		if remoteCandidate != nil {
+			a.selector.HandleBindingRequest(m, local, remoteCandidate)
+		}
 	}
 
 	if remoteCandidate != nil {
diff --git a/agent_config.go b/agent_config.go
index 93e88890..08f40128 100644
--- a/agent_config.go
+++ b/agent_config.go
@@ -200,6 +200,10 @@ type AgentConfig struct {
 	// * Implement draft-thatcher-ice-renomination
 	// * Implement custom CandidatePair switching logic
 	BindingRequestHandler func(m *stun.Message, local, remote Candidate, pair *CandidatePair) bool
+
+	// DisablePeerReflexiveDiscovery can be used to disable remote candidate discovery as peer reflexive candidate.
+	// If disabled, only added remote candidates will be used.
+	DisablePeerReflexiveDiscovery bool
 }
 
 // initWithDefaults populates an agent and falls back to defaults if fields are unset
diff --git a/agent_test.go b/agent_test.go
index 168fcddb..bbc84b4e 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -103,6 +103,48 @@ func TestHandlePeerReflexive(t *testing.T) {
 		}))
 	})
 
+	t.Run("UDP prflx candidate not discovered in handleInbound()", func(t *testing.T) {
+		var config AgentConfig
+		config.DisablePeerReflexiveDiscovery = true
+		runAgentTest(t, &config, func(ctx context.Context, a *Agent) {
+			a.selector = &controllingSelector{agent: a, log: a.log}
+
+			hostConfig := CandidateHostConfig{
+				Network:   "udp",
+				Address:   "192.168.0.2",
+				Port:      777,
+				Component: 1,
+			}
+			local, err := NewCandidateHost(&hostConfig)
+			local.conn = &fakenet.MockPacketConn{}
+			if err != nil {
+				t.Fatalf("failed to create a new candidate: %v", err)
+			}
+
+			remote := &net.UDPAddr{IP: net.ParseIP("172.17.0.3"), Port: 999}
+
+			msg, err := stun.Build(stun.BindingRequest, stun.TransactionID,
+				stun.NewUsername(a.localUfrag+":"+a.remoteUfrag),
+				UseCandidate(),
+				AttrControlling(a.tieBreaker),
+				PriorityAttr(local.Priority()),
+				stun.NewShortTermIntegrity(a.localPwd),
+				stun.Fingerprint,
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			// nolint: contextcheck
+			a.handleInbound(msg, local, remote)
+
+			// Length of remote candidate list must be zero still as peer reflexive discovery has been disabled
+			if len(a.remoteCandidates) != 0 {
+				t.Fatal("unexpeted entry in the remote candidate list")
+			}
+		})
+	})
+
 	t.Run("Bad network type with handleInbound()", func(t *testing.T) {
 		a, err := NewAgent(&AgentConfig{})
 		require.NoError(t, err)