Repository: pion/ice
Pull request #419 (2022-02-09 16:03:21)

Use UDPMux for Server reflexive candidates
#### Description
See reference issue for details https://github.com/pion/ice/issues/400

**Core changes:**
* Modified [UDPMux](https://github.com/pion/ice/blob/master/udp_mux.go) to additionally support Server Reflexive candidates (besides Host candidates).
* Added UDPMuxSrflx to Agent
* If UDPMuxSrflx is not provided, the original logic for server reflexive candidates is applied.

The reason why we went with 2 separate UDPMux'ex for Host and Srflx candidates is the conflicts between different candidate types on candidates gathering. Original UDPMux uses uFrag to mux connections which is fine for the case when there is a single candidate type. When there are more different candidate types the packets are getting mixed up. 
There is a possible solution to that which is using the STUN transaction ID map. The logic will be more complicated. 

**Next Steps:**
We will introduce a **UDPMuxRelay** alongside **UDPMux** (for host candidates) and **UDPMuxSrflx** (for server reflexive candidates). At this point, there will be 1 UDP port for each candidate type (host, srflx, relay).

We will then consolidate these 3 UDPMux'es into a single one. I believe that once all the candidate type UDPMux'es are ready we'd be able to come up with a better refactor.


#### Reference issue
https://github.com/pion/ice/issues/400

**EDIT:**

This change breaks the UDPMux interface by adding the GetXORMappedAddr function.
See https://github.com/pion/ice/pull/419#discussion_r802930982

codecov[bot]: # [Codecov](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) Report
> Merging [#419](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (868781b) into [master](https://codecov.io/gh/pion/ice/commit/4f340069615e1c5f62a857ce4fb3c3cb9bb2baaa?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) (4f34006) will **decrease** coverage by `53.87%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pion/ice/pull/419/graphs/tree.svg?width=650&height=150&src=pr&token=1ivIyHDzva&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)

```diff
@@             Coverage Diff             @@
##           master     #419       +/-   ##
===========================================
- Coverage   79.31%   25.43%   -53.88%     
===========================================
  Files          33       33               
  Lines        3650     3821      +171     
===========================================
- Hits         2895      972     -1923     
- Misses        589     2823     +2234     
+ Partials      166       26      -140     
```

| Flag | Coverage Δ | |
|---|---|---|
| go | `?` | |
| wasm | `25.43% <0.00%> (-1.20%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | Coverage Δ | |
|---|---|---|
| [agent.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-YWdlbnQuZ28=) | `0.00% <0.00%> (-80.85%)` | :arrow_down: |
| [agent\_config.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-YWdlbnRfY29uZmlnLmdv) | `0.00% <ø> (-82.15%)` | :arrow_down: |
| [gather.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-Z2F0aGVyLmdv) | `0.00% <0.00%> (-66.67%)` | :arrow_down: |
| [udp\_mux.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-dWRwX211eC5nbw==) | `0.00% <0.00%> (-81.98%)` | :arrow_down: |
| [agent\_stats.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-YWdlbnRfc3RhdHMuZ28=) | `0.00% <0.00%> (-84.38%)` | :arrow_down: |
| [selection.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-c2VsZWN0aW9uLmdv) | `0.00% <0.00%> (-82.08%)` | :arrow_down: |
| [candidate\_peer\_reflexive.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-Y2FuZGlkYXRlX3BlZXJfcmVmbGV4aXZlLmdv) | `0.00% <0.00%> (-80.65%)` | :arrow_down: |
| [transport.go](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion#diff-dHJhbnNwb3J0Lmdv) | `0.00% <0.00%> (-79.73%)` | :arrow_down: |
| ... and [15 more](https://codecov.io/gh/pion/ice/pull/419/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion). Last update [4f34006...868781b](https://codecov.io/gh/pion/ice/pull/419?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pion).

braginini: @davidzhao @Sean-Der 
Updated the PR as per suggestions of @davidzhao 

I kept **UDPMux** and **UDPMuxDefault** as is and added **UniversalUDPMux** that embeds UDPMuxDefault.
UniversalUDPMux is supposed to work with any type of candidate but so far it works only with Server Reflexive. 
I will be adding relay and host candidates support in the following PR.

The core idea of UniversalUDPMuxDefault is to just handle packets coming from the STUN server discovering a [XorMappedAddress](https://github.com/wiretrustee/ice/blob/138db20d36ad19d6561cb6f4739ca4946fea51f6/udp_mux_universal.go#L112). 
All other packets are just forwarded to the UDPMuxDefault that does the actual muxing. 
(UDPMuxDefault is embedded in the [UniversalUDPMuxDefault](https://github.com/wiretrustee/ice/blob/138db20d36ad19d6561cb6f4739ca4946fea51f6/udp_mux_universal.go#L61)).
All this is done by overriding [ReadFrom](https://github.com/wiretrustee/ice/blob/138db20d36ad19d6561cb6f4739ca4946fea51f6/udp_mux_universal.go#L90) of the original UDPConn

I used [STUN URL](https://github.com/wiretrustee/ice/blob/138db20d36ad19d6561cb6f4739ca4946fea51f6/gather.go#L370) to add some uniqueness to the muxedconn and to support multiple STUN servers. 

**Important** I had to modify the original UDPMuxDefault.RemoveConnByUfrag to avoid code duplication.
The change is pretty simple and shouldn't affect anything. Instead of removing conn by exact ufrag match I remove it by [prefix](https://github.com/wiretrustee/ice/blob/138db20d36ad19d6561cb6f4739ca4946fea51f6/udp_mux.go#L109).
This allows me to remove all the connections created by UniversalUDPMuxDefault (with the STUN URL suffix).

I tried to put as many comments as possible but let me know if you need further explanations.
What do you think guys? 
braginini: Hey @Sean-Der @davidzhao 

I reset our fork branch to a clean state and commit the whole work once again.
The reason being Github actions that I completely missed - there were many failing (e.g. lint).
I enabled actions in our fork, fixed all the issues and added more tests. All is green now :)

I'm very sorry for not doing it from the very beginning!
Looking forward for you to checking my work.

Cheers,
Misha



Patch:
diff --git a/agent.go b/agent.go
index b4aba19..aa2a1d6 100644
--- a/agent.go
+++ b/agent.go
@@ -122,9 +122,10 @@ type Agent struct {
 	loggerFactory logging.LoggerFactory
 	log           logging.LeveledLogger
 
-	net    *vnet.Net
-	tcpMux TCPMux
-	udpMux UDPMux
+	net         *vnet.Net
+	tcpMux      TCPMux
+	udpMux      UDPMux
+	udpMuxSrflx UniversalUDPMux
 
 	interfaceFilter func(string) bool
 
@@ -319,6 +320,7 @@ func NewAgent(config *AgentConfig) (*Agent, error) { //nolint:gocognit
 		a.tcpMux = newInvalidTCPMux()
 	}
 	a.udpMux = config.UDPMux
+	a.udpMuxSrflx = config.UDPMuxSrflx
 
 	if a.net == nil {
 		a.net = vnet.NewNet(nil)
@@ -892,6 +894,9 @@ func (a *Agent) removeUfragFromMux() {
 	if a.udpMux != nil {
 		a.udpMux.RemoveConnByUfrag(a.localUfrag)
 	}
+	if a.udpMuxSrflx != nil {
+		a.udpMuxSrflx.RemoveConnByUfrag(a.localUfrag)
+	}
 }
 
 // Close cleans up the Agent
diff --git a/agent_config.go b/agent_config.go
index e577af1..b373939 100644
--- a/agent_config.go
+++ b/agent_config.go
@@ -150,6 +150,12 @@ type AgentConfig struct {
 	// defer to UDPMux for incoming connections
 	UDPMux UDPMux
 
+	// UDPMuxSrflx is used for multiplexing multiple incoming UDP connections of server reflexive candidates
+	// on a single port when this is set, the agent ignores PortMin and PortMax configurations and will
+	// defer to UDPMuxSrflx for incoming connections
+	// It embeds UDPMux to do the actual connection multiplexing
+	UDPMuxSrflx UniversalUDPMux
+
 	// Proxy Dialer is a dialer that should be implemented by the user based on golang.org/x/net/proxy
 	// dial interface in order to support corporate proxies
 	ProxyDialer proxy.Dialer
diff --git a/errors.go b/errors.go
index 8ca9c2c..2ef595e 100644
--- a/errors.go
+++ b/errors.go
@@ -135,4 +135,8 @@ var (
 	errICEWriteSTUNMessage           = errors.New("the ICE conn can't write STUN messages")
 	errUDPMuxDisabled                = errors.New("UDPMux is not enabled")
 	errCandidateIPNotFound           = errors.New("could not determine local IP for Mux candidate")
+	errNoXorAddrMapping              = errors.New("no address mapping")
+	errSendSTUNPacket                = errors.New("failed to send STUN packet")
+	errXORMappedAddrTimeout          = errors.New("timeout while waiting for XORMappedAddr")
+	errNotImplemented                = errors.New("not implemented yet")
 )
diff --git a/gather.go b/gather.go
index e248c08..8d2c40d 100644
--- a/gather.go
+++ b/gather.go
@@ -97,7 +97,11 @@ func (a *Agent) gatherCandidates(ctx context.Context) {
 		case CandidateTypeServerReflexive:
 			wg.Add(1)
 			go func() {
-				a.gatherCandidatesSrflx(ctx, a.urls, a.networkTypes)
+				if a.udpMuxSrflx != nil {
+					a.gatherCandidatesSrflxUDPMux(ctx, a.urls, a.networkTypes)
+				} else {
+					a.gatherCandidatesSrflx(ctx, a.urls, a.networkTypes)
+				}
 				wg.Done()
 			}()
 			if a.extIPMapper != nil && a.extIPMapper.candidateType == CandidateTypeServerReflexive {
@@ -333,6 +337,68 @@ func (a *Agent) gatherCandidatesSrflxMapped(ctx context.Context, networkTypes []
 	}
 }
 
+func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*URL, networkTypes []NetworkType) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	for _, networkType := range networkTypes {
+		if networkType.IsTCP() {
+			continue
+		}
+
+		for i := range urls {
+			wg.Add(1)
+			go func(url URL, network string) {
+				defer wg.Done()
+
+				hostPort := fmt.Sprintf("%s:%d", url.Host, url.Port)
+				serverAddr, err := a.net.ResolveUDPAddr(network, hostPort)
+				if err != nil {
+					a.log.Warnf("failed to resolve stun host: %s: %v", hostPort, err)
+					return
+				}
+
+				xoraddr, err := a.udpMuxSrflx.GetXORMappedAddr(serverAddr, stunGatherTimeout)
+				if err != nil {
+					a.log.Warnf("could not get server reflexive address %s %s: %v\n", network, url, err)
+					return
+				}
+
+				conn, err := a.udpMuxSrflx.GetConnForURL(a.localUfrag, url.String())
+				if err != nil {
+					a.log.Warnf("could not find connection in UDPMuxSrflx %s %s: %v\n", network, url, err)
+					return
+				}
+
+				ip := xoraddr.IP
+				port := xoraddr.Port
+
+				laddr := conn.LocalAddr().(*net.UDPAddr)
+				srflxConfig := CandidateServerReflexiveConfig{
+					Network:   network,
+					Address:   ip.String(),
+					Port:      port,
+					Component: ComponentRTP,
+					RelAddr:   laddr.IP.String(),
+					RelPort:   laddr.Port,
+				}
+				c, err := NewCandidateServerReflexive(&srflxConfig)
+				if err != nil {
+					closeConnAndLog(conn, a.log, fmt.Sprintf("Failed to create server reflexive candidate: %s %s %d: %v\n", network, ip, port, 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)
+					}
+					a.log.Warnf("Failed to append to localCandidates and run onCandidateHdlr: %v\n", err)
+				}
+			}(*urls[i], networkType.String())
+		}
+	}
+}
+
 func (a *Agent) gatherCandidatesSrflx(ctx context.Context, urls []*URL, networkTypes []NetworkType) {
 	var wg sync.WaitGroup
 	defer wg.Wait()
diff --git a/go.mod b/go.mod
index 35b8113..f52c5dc 100644
--- a/go.mod
+++ b/go.mod
@@ -11,7 +11,7 @@ require (
 	github.com/pion/randutil v0.1.0
 	github.com/pion/stun v0.3.5
 	github.com/pion/transport v0.13.0
-	github.com/pion/turn/v2 v2.0.7
+	github.com/pion/turn/v2 v2.0.8
 	github.com/stretchr/testify v1.7.0
 	golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
diff --git a/go.sum b/go.sum
index 9e90feb..f6613f7 100644
--- a/go.sum
+++ b/go.sum
@@ -20,8 +20,8 @@ github.com/pion/stun v0.3.5/go.mod h1:gDMim+47EeEtfWogA37n6qXZS88L5V6LqFcf+DZA2U
 github.com/pion/transport v0.12.2/go.mod h1:N3+vZQD9HlDP5GWkZ85LohxNsDcNgofQmyL6ojX5d8Q=
 github.com/pion/transport v0.13.0 h1:KWTA5ZrQogizzYwPEciGtHPLwpAjE91FgXnyu+Hv2uY=
 github.com/pion/transport v0.13.0/go.mod h1:yxm9uXpK9bpBBWkITk13cLo1y5/ur5VQpG22ny6EP7g=
-github.com/pion/turn/v2 v2.0.7 h1:SZhc00WDovK6czaN1RSiHqbwANtIO6wfZQsU0m0KNE8=
-github.com/pion/turn/v2 v2.0.7/go.mod h1:+y7xl719J8bAEVpSXBXvTxStjJv3hbz9YFflvkpcGPw=
+github.com/pion/turn/v2 v2.0.8 h1:KEstL92OUN3k5k8qxsXHpr7WWfrdp7iJZHx99ud8muw=
+github.com/pion/turn/v2 v2.0.8/go.mod h1:+y7xl719J8bAEVpSXBXvTxStjJv3hbz9YFflvkpcGPw=
 github.com/pion/udp v0.1.1 h1:8UAPvyqmsxK8oOjloDk4wUt63TzFe9WEJkg5lChlj7o=
 github.com/pion/udp v0.1.1/go.mod h1:6AFo+CMdKQm7UiA0eUPA8/eVCTx8jBIITLZHc9DWX5M=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
diff --git a/udp_mux_universal.go b/udp_mux_universal.go
new file mode 100644
index 0000000..c619877
--- /dev/null
+++ b/udp_mux_universal.go
@@ -0,0 +1,265 @@
+package ice
+
+import (
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/pion/logging"
+	"github.com/pion/stun"
+)
+
+// UniversalUDPMux allows multiple connections to go over a single UDP port for
+// host, server reflexive and relayed candidates.
+// Actual connection muxing is happening in the UDPMux.
+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) (net.PacketConn, error)
+}
+
+// UniversalUDPMuxDefault handles STUN and TURN servers packets by wrapping the original UDPConn overriding ReadFrom.
+// It the passes packets to the UDPMux that does the actual connection muxing.
+type UniversalUDPMuxDefault struct {
+	*UDPMuxDefault
+	params UniversalUDPMuxParams
+
+	// since we have a shared socket, for srflx candidates it makes sense to have a shared mapped address across all the agents
+	// stun.XORMappedAddress indexed by the STUN server addr
+	xorMappedMap map[string]*xorMapped
+}
+
+// UniversalUDPMuxParams are parameters for UniversalUDPMux server reflexive.
+type UniversalUDPMuxParams struct {
+	Logger                logging.LeveledLogger
+	UDPConn               net.PacketConn
+	XORMappedAddrCacheTTL time.Duration
+}
+
+// NewUniversalUDPMuxDefault creates an implementation of UniversalUDPMux embedding UDPMux
+func NewUniversalUDPMuxDefault(params UniversalUDPMuxParams) *UniversalUDPMuxDefault {
+	if params.Logger == nil {
+		params.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
+	}
+	if params.XORMappedAddrCacheTTL == 0 {
+		params.XORMappedAddrCacheTTL = time.Second * 25
+	}
+
+	m := &UniversalUDPMuxDefault{
+		params:       params,
+		xorMappedMap: make(map[string]*xorMapped),
+	}
+
+	// wrap UDP connection, process server reflexive messages
+	// before they are passed to the UDPMux connection handler (connWorker)
+	m.params.UDPConn = &udpConn{
+		PacketConn: params.UDPConn,
+		mux:        m,
+		logger:     params.Logger,
+	}
+
+	// embed UDPMux
+	udpMuxParams := UDPMuxParams{
+		Logger:  params.Logger,
+		UDPConn: m.params.UDPConn,
+	}
+	m.UDPMuxDefault = NewUDPMuxDefault(udpMuxParams)
+
+	return m
+}
+
+// udpConn is a wrapper around UDPMux conn that overrides ReadFrom and handles STUN/TURN packets
+type udpConn struct {
+	net.PacketConn
+	mux    *UniversalUDPMuxDefault
+	logger logging.LeveledLogger
+}
+
+// GetRelayedAddr creates relayed connection to the given TURN service and returns the relayed addr.
+// Not implemented yet.
+func (m *UniversalUDPMuxDefault) GetRelayedAddr(turnAddr net.Addr, deadline time.Duration) (*net.Addr, error) {
+	return nil, errNotImplemented
+}
+
+// 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) (net.PacketConn, error) {
+	return m.UDPMuxDefault.GetConn(fmt.Sprintf("%s%s", ufrag, url))
+}
+
+// ReadFrom is called by UDPMux connWorker and handles packets coming from the STUN server discovering a mapped address.
+// It passes processed packets further to the UDPMux (maybe this is not really necessary).
+func (c *udpConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
+	n, addr, err = c.PacketConn.ReadFrom(p)
+	if err != nil {
+		return
+	}
+
+	if stun.IsMessage(p[:n]) {
+		msg := &stun.Message{
+			Raw: append([]byte{}, p[:n]...),
+		}
+
+		if err = msg.Decode(); err != nil {
+			c.logger.Warnf("Failed to handle decode ICE from %s: %v\n", addr.String(), err)
+			return n, addr, nil
+		}
+
+		udpAddr, ok := addr.(*net.UDPAddr)
+		if !ok {
+			// message about this err will be logged in the UDPMux
+			return
+		}
+
+		if c.mux.isXORMappedResponse(msg, udpAddr.String()) {
+			err = c.mux.handleXORMappedResponse(udpAddr, msg)
+			if err != nil {
+				c.logger.Debugf("%w: %v", errGetXorMappedAddrResponse, err)
+				return n, addr, nil
+			}
+			return
+		}
+	}
+	return n, addr, err
+}
+
+// isXORMappedResponse indicates whether the message is a XORMappedAddress and is coming from the known STUN server.
+func (m *UniversalUDPMuxDefault) isXORMappedResponse(msg *stun.Message, stunAddr string) bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	// check first if it is a STUN server address because remote peer can also send similar messages but as a BindingSuccess
+	_, ok := m.xorMappedMap[stunAddr]
+	_, err := msg.Get(stun.AttrXORMappedAddress)
+	return err == nil && ok
+}
+
+// handleXORMappedResponse parses response from the STUN server, extracts XORMappedAddress attribute
+// and set the mapped address for the server
+func (m *UniversalUDPMuxDefault) handleXORMappedResponse(stunAddr *net.UDPAddr, msg *stun.Message) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	mappedAddr, ok := m.xorMappedMap[stunAddr.String()]
+	if !ok {
+		return errNoXorAddrMapping
+	}
+
+	var addr stun.XORMappedAddress
+	if err := addr.GetFrom(msg); err != nil {
+		return err
+	}
+
+	m.xorMappedMap[stunAddr.String()] = mappedAddr
+	mappedAddr.SetAddr(&addr)
+
+	return nil
+}
+
+// GetXORMappedAddr returns *stun.XORMappedAddress if already present for a given STUN server.
+// Makes a STUN binding request to discover mapped address otherwise.
+// Blocks until the stun.XORMappedAddress has been discovered or deadline.
+// Method is safe for concurrent use.
+func (m *UniversalUDPMuxDefault) GetXORMappedAddr(serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
+	m.mu.Lock()
+	mappedAddr, ok := m.xorMappedMap[serverAddr.String()]
+	// if we already have a mapping for this STUN server (address already received)
+	// and if it is not too old we return it without making a new request to STUN server
+	if ok {
+		if mappedAddr.expired() {
+			mappedAddr.closeWaiters()
+			delete(m.xorMappedMap, serverAddr.String())
+			ok = false
+		} else if mappedAddr.pending() {
+			ok = false
+		}
+	}
+	m.mu.Unlock()
+	if ok {
+		return mappedAddr.addr, nil
+	}
+
+	// otherwise, make a STUN request to discover the address
+	// or wait for already sent request to complete
+	waitAddrReceived, err := m.sendStun(serverAddr)
+	if err != nil {
+		return nil, errSendSTUNPacket
+	}
+
+	// block until response was handled by the connWorker routine and XORMappedAddress was updated
+	select {
+	case <-waitAddrReceived:
+		// when channel closed, addr was obtained
+		m.mu.Lock()
+		mappedAddr := *m.xorMappedMap[serverAddr.String()]
+		m.mu.Unlock()
+		if mappedAddr.addr == nil {
+			return nil, errNoXorAddrMapping
+		}
+		return mappedAddr.addr, nil
+	case <-time.After(deadline):
+		return nil, errXORMappedAddrTimeout
+	}
+}
+
+// sendStun sends a STUN request via UDP conn.
+//
+// The returned channel is closed when the STUN response has been received.
+// Method is safe for concurrent use.
+func (m *UniversalUDPMuxDefault) sendStun(serverAddr net.Addr) (chan struct{}, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	// if record present in the map, we already sent a STUN request,
+	// just wait when waitAddrReceived will be closed
+	addrMap, ok := m.xorMappedMap[serverAddr.String()]
+	if !ok {
+		addrMap = &xorMapped{
+			expiresAt:        time.Now().Add(m.params.XORMappedAddrCacheTTL),
+			waitAddrReceived: make(chan struct{}),
+		}
+		m.xorMappedMap[serverAddr.String()] = addrMap
+	}
+
+	req, err := stun.Build(stun.BindingRequest, stun.TransactionID)
+	if err != nil {
+		return nil, err
+	}
+
+	if _, err = m.params.UDPConn.WriteTo(req.Raw, serverAddr); err != nil {
+		return nil, err
+	}
+
+	return addrMap.waitAddrReceived, nil
+}
+
+type xorMapped struct {
+	addr             *stun.XORMappedAddress
+	waitAddrReceived chan struct{}
+	expiresAt        time.Time
+}
+
+func (a *xorMapped) closeWaiters() {
+	select {
+	case <-a.waitAddrReceived:
+		// notify was close, ok, that means we received duplicate response
+		// just exit
+		break
+	default:
+		// notify tha twe have a new addr
+		close(a.waitAddrReceived)
+	}
+}
+
+func (a *xorMapped) pending() bool {
+	return a.addr == nil
+}
+
+func (a *xorMapped) expired() bool {
+	return a.expiresAt.Before(time.Now())
+}
+
+func (a *xorMapped) SetAddr(addr *stun.XORMappedAddress) {
+	a.addr = addr
+	a.closeWaiters()
+}


Test patch:
diff --git a/gather_test.go b/gather_test.go
index e69d987..23c8296 100644
--- a/gather_test.go
+++ b/gather_test.go
@@ -11,12 +11,14 @@ import (
 	"reflect"
 	"sort"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/pion/dtls/v2"
 	"github.com/pion/dtls/v2/pkg/crypto/selfsign"
 	"github.com/pion/logging"
+	"github.com/pion/stun"
 	"github.com/pion/transport/test"
 	"github.com/pion/turn/v2"
 	"github.com/stretchr/testify/assert"
@@ -484,3 +486,93 @@ func TestTURNProxyDialer(t *testing.T) {
 
 	assert.NoError(t, a.Close())
 }
+
+// Assert that UniversalUDPMux is used while gathering when configured in the Agent
+func TestUniversalUDPMuxUsage(t *testing.T) {
+	report := test.CheckRoutines(t)
+	defer report()
+
+	lim := test.TimeOut(time.Second * 30)
+	defer lim.Stop()
+
+	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IP{127, 0, 0, 1}, Port: randomPort(t)})
+	assert.NoError(t, err)
+	defer func() {
+		_ = conn.Close()
+	}()
+
+	udpMuxSrflx := &universalUDPMuxMock{
+		conn: conn,
+	}
+
+	numSTUNS := 3
+	urls := []*URL{}
+	for i := 0; i < numSTUNS; i++ {
+		urls = append(urls, &URL{
+			Scheme: SchemeTypeSTUN,
+			Host:   "127.0.0.1",
+			Port:   3478 + i,
+		})
+	}
+
+	a, err := NewAgent(&AgentConfig{
+		NetworkTypes:   supportedNetworkTypes(),
+		Urls:           urls,
+		CandidateTypes: []CandidateType{CandidateTypeServerReflexive},
+		UDPMuxSrflx:    udpMuxSrflx,
+	})
+	assert.NoError(t, err)
+
+	candidateGathered, candidateGatheredFunc := context.WithCancel(context.Background())
+	assert.NoError(t, a.OnCandidate(func(c Candidate) {
+		if c == nil {
+			candidateGatheredFunc()
+			return
+		}
+		t.Log(c.NetworkType(), c.Priority(), c)
+	}))
+	assert.NoError(t, a.GatherCandidates())
+
+	<-candidateGathered.Done()
+
+	assert.NoError(t, a.Close())
+	// twice because of 2 STUN servers configured
+	assert.Equal(t, numSTUNS, udpMuxSrflx.getXORMappedAddrUsedTimes, "expected times that GetXORMappedAddr should be called")
+	// one for Restart() when agent has been initialized and one time when Close() the agent
+	assert.Equal(t, 2, udpMuxSrflx.removeConnByUfragTimes, "expected times that RemoveConnByUfrag should be called")
+	// twice because of 2 STUN servers configured
+	assert.Equal(t, numSTUNS, udpMuxSrflx.getConnForURLTimes, "expected times that GetConnForURL should be called")
+}
+
+type universalUDPMuxMock struct {
+	UDPMux
+	getXORMappedAddrUsedTimes int
+	removeConnByUfragTimes    int
+	getConnForURLTimes        int
+	mu                        sync.Mutex
+	conn                      *net.UDPConn
+}
+
+func (m *universalUDPMuxMock) GetRelayedAddr(turnAddr net.Addr, deadline time.Duration) (*net.Addr, error) {
+	return nil, errNotImplemented
+}
+
+func (m *universalUDPMuxMock) GetConnForURL(ufrag string, url string) (net.PacketConn, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.getConnForURLTimes++
+	return m.conn, nil
+}
+
+func (m *universalUDPMuxMock) GetXORMappedAddr(serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.getXORMappedAddrUsedTimes++
+	return &stun.XORMappedAddress{IP: net.IP{100, 64, 0, 1}, Port: 77878}, nil
+}
+
+func (m *universalUDPMuxMock) RemoveConnByUfrag(ufrag string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.removeConnByUfragTimes++
+}
diff --git a/udp_mux_universal_test.go b/udp_mux_universal_test.go
new file mode 100644
index 0000000..c263bf6
--- /dev/null
+++ b/udp_mux_universal_test.go
@@ -0,0 +1,127 @@
+// +build !js
+
+package ice
+
+import (
+	"net"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/pion/stun"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUniversalUDPMux(t *testing.T) {
+	conn, err := net.ListenUDP(udp, &net.UDPAddr{})
+	require.NoError(t, err)
+
+	udpMux := NewUniversalUDPMuxDefault(UniversalUDPMuxParams{
+		Logger:  nil,
+		UDPConn: conn,
+	})
+
+	require.NoError(t, err)
+	defer func() {
+		_ = udpMux.Close()
+		_ = conn.Close()
+	}()
+
+	require.NotNil(t, udpMux.LocalAddr(), "tcpMux.LocalAddr() is nil")
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		testMuxSrflxConnection(t, udpMux, "ufrag4", udp)
+	}()
+
+	wg.Wait()
+}
+
+func testMuxSrflxConnection(t *testing.T, udpMux *UniversalUDPMuxDefault, ufrag string, network string) {
+	pktConn, err := udpMux.GetConn(ufrag)
+	require.NoError(t, err, "error retrieving muxed connection for ufrag")
+	defer func() {
+		_ = pktConn.Close()
+	}()
+
+	remoteConn, err := net.DialUDP(network, nil, &net.UDPAddr{
+		Port: udpMux.LocalAddr().(*net.UDPAddr).Port,
+	})
+	require.NoError(t, err, "error dialing test udp connection")
+	defer func() {
+		_ = remoteConn.Close()
+	}()
+
+	// use small value for TTL to check expiration of the address
+	udpMux.params.XORMappedAddrCacheTTL = time.Millisecond * 20
+	testXORIP := net.ParseIP("213.141.156.236")
+	testXORPort := 21254
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		address, e := udpMux.GetXORMappedAddr(remoteConn.LocalAddr(), time.Second)
+		require.NoError(t, e)
+		require.NotNil(t, address)
+		require.True(t, address.IP.Equal(testXORIP))
+		require.Equal(t, address.Port, testXORPort)
+	}()
+
+	// wait until GetXORMappedAddr calls sendStun method
+	time.Sleep(time.Millisecond)
+
+	// check that mapped address filled correctly after sent stun
+	udpMux.mu.Lock()
+	mappedAddr, ok := udpMux.xorMappedMap[remoteConn.LocalAddr().String()]
+	require.True(t, ok)
+	require.NotNil(t, mappedAddr)
+	require.True(t, mappedAddr.pending())
+	require.False(t, mappedAddr.expired())
+	udpMux.mu.Unlock()
+
+	// clean receiver read buffer
+	buf := make([]byte, receiveMTU)
+	_, err = remoteConn.Read(buf)
+	require.NoError(t, err)
+
+	// write back to udpMux XOR message with address
+	msg := stun.New()
+	msg.Type = stun.MessageType{Method: stun.MethodBinding, Class: stun.ClassRequest}
+	msg.Add(stun.AttrUsername, []byte(ufrag+":otherufrag"))
+	addr := &stun.XORMappedAddress{
+		IP:   testXORIP,
+		Port: testXORPort,
+	}
+	err = addr.AddTo(msg)
+	require.NoError(t, err)
+
+	msg.Encode()
+	_, err = remoteConn.Write(msg.Raw)
+	require.NoError(t, err)
+
+	// wait for the packet to be consumed and parsed by udpMux
+	wg.Wait()
+
+	// we should get address immediately from the cached map
+	address, err := udpMux.GetXORMappedAddr(remoteConn.LocalAddr(), time.Second)
+	require.NoError(t, err)
+	require.NotNil(t, address)
+
+	udpMux.mu.Lock()
+	// check mappedAddr is not pending, we didn't send stun twice
+	require.False(t, mappedAddr.pending())
+
+	// check expiration by TTL
+	time.Sleep(time.Millisecond * 21)
+	require.True(t, mappedAddr.expired())
+	udpMux.mu.Unlock()
+
+	// after expire, we send stun request again
+	// but we not receive response in 5 milliseconds and should get error here
+	address, err = udpMux.GetXORMappedAddr(remoteConn.LocalAddr(), time.Millisecond*5)
+	require.NotNil(t, err)
+	require.Nil(t, address)
+}