Refactore/v4 stun routing (#5)

agent.go

@@ -818,6 +818,11 @@ func (a *Agent) setCandidateExtensions(cand Candidate) {
 	if err != nil {
 		a.log.Errorf("Failed to add ufrag extension to candidate: %v", err)
 	}
+
+	err = cand.AddExtension(newCandidateIDExtension(cand.ID()))
+	if err != nil {
+		a.log.Errorf("Failed to add candidate id extension to candidate: %v", err)
+	}
 }
 
 // GetRemoteCandidates returns the remote candidates.
@@ -1153,6 +1158,14 @@ func (a *Agent) handleInbound(msg *stun.Message, local Candidate, remote net.Add
 				RelPort:   0,
 			}
 
+			// If the remote candidate has CandidatePairID, we can use it to set the CandidateID
+			candidatePairID, ok, err := CandidatePairIDFromSTUN(msg)
+			if err != nil {
+				a.log.Errorf("Failed to create candidate pair ID from STUN message (%s)", err)
+			} else if ok {
+				prflxCandidateConfig.CandidateID = candidatePairID.SourceCandidateID()
+			}
+
 			prflxCandidate, err := NewCandidatePeerReflexive(&prflxCandidateConfig)
 			if err != nil {
 				a.log.Errorf("Failed to create new remote prflx candidate (%s)", err)

candidate_base.go

@@ -434,12 +434,13 @@ func (c *candidateBase) DeepEqual(other Candidate) bool {
 // String makes the candidateBase printable.
 func (c *candidateBase) String() string {
 	return fmt.Sprintf(
-		"%s %s %s%s (resolved: %v)",
+		"%s %s %s%s (resolved: %v) %s",
 		c.NetworkType(),
 		c.Type(),
 		net.JoinHostPort(c.Address(), strconv.Itoa(c.Port())),
 		c.relatedAddress,
 		c.resolvedAddr,
+		c.id,
 	)
 }
 
@@ -791,12 +792,14 @@ func UnmarshalCandidate(raw string) (Candidate, error) { //nolint:cyclop
 		}
 	}
 
+	candidateID := candidateIDFromExtensions(extensions)
+
 	// this code is ugly because we can't break backwards compatibility
 	// with the old way of parsing candidates
 	switch typ {
 	case "host":
 		candidate, err := NewCandidateHost(&CandidateHostConfig{
-			"",
+			candidateID,
 			protocol,
 			address,
 			port,
@@ -815,7 +818,7 @@ func UnmarshalCandidate(raw string) (Candidate, error) { //nolint:cyclop
 		return candidate, nil
 	case "srflx":
 		candidate, err := NewCandidateServerReflexive(&CandidateServerReflexiveConfig{
-			"",
+			candidateID,
 			protocol,
 			address,
 			port,
@@ -834,7 +837,7 @@ func UnmarshalCandidate(raw string) (Candidate, error) { //nolint:cyclop
 		return candidate, nil
 	case "prflx":
 		candidate, err := NewCandidatePeerReflexive(&CandidatePeerReflexiveConfig{
-			"",
+			candidateID,
 			protocol,
 			address,
 			port,
@@ -853,7 +856,7 @@ func UnmarshalCandidate(raw string) (Candidate, error) { //nolint:cyclop
 		return candidate, nil
 	case "relay":
 		candidate, err := NewCandidateRelay(&CandidateRelayConfig{
-			"",
+			candidateID,
 			protocol,
 			address,
 			port,

candidatepair.go

@@ -125,6 +125,9 @@ func (p *CandidatePair) Write(b []byte) (int, error) {
 }
 
 func (a *Agent) sendSTUN(msg *stun.Message, local, remote Candidate) {
+	// Add the custom attribute to the message
+	msg.Add(AttrCandidatePairID, []byte(NewCandidatePairID(local, remote).String()))
+
 	_, err := local.writeTo(msg.Raw, remote)
 	if err != nil {
 		a.log.Tracef("Failed to send STUN message: %s", err)

extension_candidate_pair_id.go

@@ -0,0 +1,30 @@
+package ice
+
+import (
+	"fmt"
+	"strings"
+)
+
+const (
+	// used in RDP for candidate ID extension
+	ExtensionKeyCandidateID = "cid"
+
+	candidateIDPrefix = "candidate:"
+)
+
+func candidateIDFromExtensions(extensions []CandidateExtension) string {
+	for _, ext := range extensions {
+		if ext.Key == ExtensionKeyCandidateID {
+			return fmt.Sprintf("candidate:%s", ext.Value)
+		}
+	}
+
+	return ""
+}
+
+func newCandidateIDExtension(candidateID string) CandidateExtension {
+	return CandidateExtension{
+		Key:   ExtensionKeyCandidateID,
+		Value: strings.TrimPrefix(candidateID, candidateIDPrefix),
+	}
+}

gather.go

@@ -383,18 +383,17 @@ func (a *Agent) gatherCandidatesLocalUDPMux(ctx context.Context) error { //nolin
 			continue
 		}
 
-		conn, err := a.udpMux.GetConn(a.localUfrag, udpAddr)
-		if err != nil {
-			return err
-		}
-
 		c, err := NewCandidateHost(&hostConfig)
 		if err != nil {
-			closeConnAndLog(conn, a.log, "failed to create host mux candidate: %s %d: %v", candidateIP, udpAddr.Port, err)
-
+			closeConnAndLog(nil, a.log, "failed to create host mux candidate: %s %d: %v", candidateIP, udpAddr.Port, err)
 			continue
 		}
 
+		conn, err := a.udpMux.GetConn(a.localUfrag, udpAddr, c.ID())
+		if err != nil {
+			return err
+		}
+
 		if err := a.addCandidate(ctx, c, conn); err != nil {
 			if closeErr := c.close(); closeErr != nil {
 				a.log.Warnf("Failed to close candidate: %v", closeErr)
@@ -531,13 +530,6 @@ func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*stun.UR
 						return
 					}
 
-					conn, err := a.udpMuxSrflx.GetConnForURL(a.localUfrag, url.String(), localAddr)
-					if err != nil {
-						a.log.Warnf("Failed to find connection in UDPMuxSrflx %s %s: %v", network, url, err)
-
-						return
-					}
-
 					ip := xorAddr.IP
 					port := xorAddr.Port
 
@@ -551,11 +543,17 @@ func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*stun.UR
 					}
 					c, err := NewCandidateServerReflexive(&srflxConfig)
 					if err != nil {
-						closeConnAndLog(conn, a.log, "failed to create server reflexive candidate: %s %s %d: %v", network, ip, port, err)
+						closeConnAndLog(nil, a.log, "failed to create server reflexive candidate: %s %s %d: %v", network, ip, port, err)
 
 						return
 					}
 
+					conn, err := a.udpMuxSrflx.GetConnForURL(a.localUfrag, url.String(), localAddr, c.ID())
+					if err != nil {
+						a.log.Warnf("Failed to find connection in UDPMuxSrflx %s %s: %v", network, url, err)
+						return
+					}
+
 					if err := a.addCandidate(ctx, c, conn); err != nil {
 						if closeErr := c.close(); closeErr != nil {
 							a.log.Warnf("Failed to close candidate: %v", closeErr)

gather_test.go

@@ -888,7 +888,7 @@ func (m *universalUDPMuxMock) GetRelayedAddr(net.Addr, time.Duration) (*net.Addr
 	return nil, errNotImplemented
 }
 
-func (m *universalUDPMuxMock) GetConnForURL(string, string, net.Addr) (net.PacketConn, error) {
+func (m *universalUDPMuxMock) GetConnForURL(string, string, net.Addr, string) (net.PacketConn, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	m.getConnForURLTimes++

stun_addon.go

@@ -0,0 +1,64 @@
+package ice
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/pion/stun/v3"
+)
+
+const (
+	AttrCandidatePairID stun.AttrType = 0x8100 // Custom attribute for candidate ID
+)
+
+type CandidatePairID struct {
+	source      string
+	destination string
+}
+
+func CandidatePairIDFromSTUN(msg *stun.Message) (*CandidatePairID, bool, error) {
+	candidatePairIDBytes, err := msg.Get(AttrCandidatePairID)
+	if err != nil {
+		return nil, false, nil
+	}
+	candidatePairID, err := ParseCandidatePairID(candidatePairIDBytes)
+	if err != nil {
+		return nil, false, err
+	}
+	return &candidatePairID, true, nil
+}
+
+// NewCandidatePairID creates a CandidatePairID from local and remote candidates.
+func NewCandidatePairID(local, remote Candidate) CandidatePairID {
+	localID := strings.TrimPrefix(local.ID(), "candidate:")
+	remoteID := strings.TrimPrefix(remote.ID(), "candidate:")
+
+	return CandidatePairID{
+		source:      localID,
+		destination: remoteID,
+	}
+}
+
+// ParseCandidatePairID parses a CandidatePairID from its string representation.
+func ParseCandidatePairID(id []byte) (CandidatePairID, error) {
+	parts := strings.SplitN(string(id), ":", 2)
+	if len(parts) != 2 {
+		return CandidatePairID{}, fmt.Errorf("invalid candidates ID format: %s", id)
+	}
+	return CandidatePairID{
+		source:      parts[0],
+		destination: parts[1],
+	}, nil
+}
+
+func (cp CandidatePairID) String() string {
+	return fmt.Sprintf("%s:%s", cp.source, cp.destination)
+}
+
+func (cp CandidatePairID) SourceCandidateID() string {
+	return fmt.Sprintf("candidate:%s", cp.source)
+}
+
+func (cp CandidatePairID) TargetCandidateID() string {
+	return fmt.Sprintf("candidate:%s", cp.destination)
+}

udp_mux.go

@@ -21,7 +21,7 @@ import (
 // UDPMux allows multiple connections to go over a single UDP port.
 type UDPMux interface {
 	io.Closer
-	GetConn(ufrag string, addr net.Addr) (net.PacketConn, error)
+	GetConn(ufrag string, addr net.Addr, candidateID string) (net.PacketConn, error)
 	RemoveConnByUfrag(ufrag string)
 	GetListenAddresses() []net.Addr
 }
@@ -146,7 +146,7 @@ func (m *UDPMuxDefault) GetListenAddresses() []net.Addr {
 
 // GetConn returns a PacketConn given the connection's ufrag and network address.
 // creates the connection if an existing one can't be found.
-func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, error) {
+func (m *UDPMuxDefault) GetConn(ufrag string, addr net.Addr, _ string) (net.PacketConn, error) {
 	// don't check addr for mux using unspecified address
 	if len(m.localAddrsForUnspecified) == 0 && m.params.UDPConnString != addr.String() {
 		return nil, errInvalidAddress

udp_mux_multi.go

@@ -38,13 +38,13 @@ func NewMultiUDPMuxDefault(muxes ...UDPMux) *MultiUDPMuxDefault {
 
 // GetConn returns a PacketConn given the connection's ufrag and network
 // creates the connection if an existing one can't be found.
-func (m *MultiUDPMuxDefault) GetConn(ufrag string, addr net.Addr) (net.PacketConn, error) {
+func (m *MultiUDPMuxDefault) GetConn(ufrag string, addr net.Addr, candidateID string) (net.PacketConn, error) {
 	mux, ok := m.localAddrToMux[addr.String()]
 	if !ok {
 		return nil, errNoUDPMuxAvailable
 	}
 
-	return mux.GetConn(ufrag, addr)
+	return mux.GetConn(ufrag, addr, candidateID)
 }
 
 // RemoveConnByUfrag stops and removes the muxed packet connection

udp_mux_multi_test.go

@@ -74,7 +74,7 @@ func TestMultiUDPMux(t *testing.T) {
 	require.NoError(t, udpMuxMulti.Close())
 
 	// Can't create more connections
-	_, err = udpMuxMulti.GetConn("failufrag", conn1.LocalAddr())
+	_, err = udpMuxMulti.GetConn("failufrag", conn1.LocalAddr(), "")
 	require.Error(t, err)
 }
 
@@ -91,7 +91,7 @@ func testMultiUDPMuxConnections(t *testing.T, udpMuxMulti *MultiUDPMuxDefault, u
 		} else if network == udp6 && udpAddr.IP.To4() != nil {
 			continue
 		}
-		c, err := udpMuxMulti.GetConn(ufrag, addr)
+		c, err := udpMuxMulti.GetConn(ufrag, addr, "")
 		require.NoError(t, err, "error retrieving muxed connection for ufrag")
 		pktConns = append(pktConns, c)
 	}

udp_mux_universal.go

@@ -20,7 +20,7 @@ type UniversalUDPMux interface {
 	UDPMux
 	GetXORMappedAddr(stunAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error)
 	GetRelayedAddr(turnAddr net.Addr, deadline time.Duration) (*net.Addr, error)
-	GetConnForURL(ufrag string, url string, addr net.Addr) (net.PacketConn, error)
+	GetConnForURL(ufrag string, url string, addr net.Addr, candidatedID string) (net.PacketConn, error)
 }
 
 // UniversalUDPMuxDefault handles STUN and TURN servers packets by wrapping the original UDPConn overriding ReadFrom.
@@ -92,8 +92,8 @@ func (m *UniversalUDPMuxDefault) GetRelayedAddr(net.Addr, time.Duration) (*net.A
 // GetConnForURL add uniques to the muxed connection by concatenating ufrag and URL
 // (e.g. STUN URL) to be able to support multiple STUN/TURN servers
 // and return a unique connection per server.
-func (m *UniversalUDPMuxDefault) GetConnForURL(ufrag string, url string, addr net.Addr) (net.PacketConn, error) {
-	return m.UDPMuxDefault.GetConn(fmt.Sprintf("%s%s", ufrag, url), addr)
+func (m *UniversalUDPMuxDefault) GetConnForURL(ufrag string, url string, addr net.Addr, candidateID string) (net.PacketConn, error) {
+	return m.UDPMuxDefault.GetConn(fmt.Sprintf("%s%s", ufrag, url), addr, candidateID)
 }
 
 // ReadFrom is called by UDPMux connWorker and handles packets coming from the STUN server discovering a mapped address.