From d21ae5e0e54465db200d27e6c61d2f44eb012aad Mon Sep 17 00:00:00 2001
From: Joe Turki <git@joeturki.com>
Date: Thu, 30 Jan 2025 01:14:01 -0600
Subject: [PATCH] Include ufrag in generated ICE candidates

Include ufrag extension in the ICE candidates generated by the ICE agent
---
 agent.go      | 11 +++++++++++
 agent_test.go | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/agent.go b/agent.go
index 4ce1b1f9..fd1faaea 100644
--- a/agent.go
+++ b/agent.go
@@ -799,6 +799,7 @@ func (a *Agent) addCandidate(ctx context.Context, cand Candidate, candidateConn
 			}
 		}
 
+		a.setCandidateExtensions(cand)
 		cand.start(a, candidateConn, a.startedCh)
 
 		set = append(set, cand)
@@ -818,6 +819,16 @@ func (a *Agent) addCandidate(ctx context.Context, cand Candidate, candidateConn
 	})
 }
 
+func (a *Agent) setCandidateExtensions(cand Candidate) {
+	err := cand.AddExtension(CandidateExtension{
+		Key:   "ufrag",
+		Value: a.localUfrag,
+	})
+	if err != nil {
+		a.log.Errorf("Failed to add ufrag extension to candidate: %v", err)
+	}
+}
+
 // GetRemoteCandidates returns the remote candidates.
 func (a *Agent) GetRemoteCandidates() ([]Candidate, error) {
 	var res []Candidate
diff --git a/agent_test.go b/agent_test.go
index f1d7f36f..42a18465 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -2071,3 +2071,42 @@ func TestAgentGracefulCloseDeadlock(t *testing.T) {
 	closeNow.Done()
 	closed.Wait()
 }
+
+func TestSetCandidatesUfrag(t *testing.T) {
+	var config AgentConfig
+
+	agent, err := NewAgent(&config)
+	if err != nil {
+		t.Fatalf("Error constructing ice.Agent: %v", err)
+	}
+	defer func() {
+		require.NoError(t, agent.Close())
+	}()
+
+	dummyConn := &net.UDPConn{}
+
+	for i := 0; i < 5; i++ {
+		cfg := CandidateHostConfig{
+			Network:   "udp",
+			Address:   "192.168.0.2",
+			Port:      1000 + i,
+			Component: 1,
+		}
+
+		cand, errCand := NewCandidateHost(&cfg)
+		require.NoError(t, errCand)
+
+		err = agent.addCandidate(context.Background(), cand, dummyConn)
+		require.NoError(t, err)
+	}
+
+	actualCandidates, err := agent.GetLocalCandidates()
+	require.NoError(t, err)
+
+	for _, candidate := range actualCandidates {
+		ext, ok := candidate.GetExtension("ufrag")
+
+		require.True(t, ok)
+		require.Equal(t, agent.localUfrag, ext.Value)
+	}
+}