From 6202be846b513e053ce5eaaff63ed890728855be Mon Sep 17 00:00:00 2001 From: Zoltan Papp Date: Mon, 8 Sep 2025 20:49:34 +0200 Subject: [PATCH] Refactore/v4 stun routing (#5) --- agent.go | 13 +++++++ candidate_base.go | 13 ++++--- candidatepair.go | 3 ++ extension_candidate_pair_id.go | 30 ++++++++++++++++ gather.go | 28 +++++++-------- gather_test.go | 2 +- stun_addon.go | 64 ++++++++++++++++++++++++++++++++++ udp_mux.go | 4 +-- udp_mux_multi.go | 4 +-- udp_mux_multi_test.go | 4 +-- udp_mux_universal.go | 6 ++-- 11 files changed, 141 insertions(+), 30 deletions(-) create mode 100644 extension_candidate_pair_id.go create mode 100644 stun_addon.go diff --git a/agent.go b/agent.go index 5220c1d..765d03b 100644 --- a/agent.go +++ b/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) diff --git a/candidate_base.go b/candidate_base.go index 45c090e..cc88dd5 100644 --- a/candidate_base.go +++ b/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, diff --git a/candidatepair.go b/candidatepair.go index 82655aa..a1537af 100644 --- a/candidatepair.go +++ b/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) diff --git a/extension_candidate_pair_id.go b/extension_candidate_pair_id.go new file mode 100644 index 0000000..2f00e9d --- /dev/null +++ b/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), + } +} diff --git a/gather.go b/gather.go index ebf2999..0a87f4d 100644 --- a/gather.go +++ b/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) diff --git a/gather_test.go b/gather_test.go index 6ef39e1..425a915 100644 --- a/gather_test.go +++ b/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++ diff --git a/stun_addon.go b/stun_addon.go new file mode 100644 index 0000000..259fb88 --- /dev/null +++ b/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) +} diff --git a/udp_mux.go b/udp_mux.go index 257ef5f..c38c3b3 100644 --- a/udp_mux.go +++ b/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 diff --git a/udp_mux_multi.go b/udp_mux_multi.go index 46c88bb..44d0523 100644 --- a/udp_mux_multi.go +++ b/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 diff --git a/udp_mux_multi_test.go b/udp_mux_multi_test.go index 4fc64de..f236607 100644 --- a/udp_mux_multi_test.go +++ b/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) } diff --git a/udp_mux_universal.go b/udp_mux_universal.go index 50c80af..9a17e0d 100644 --- a/udp_mux_universal.go +++ b/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.