Fix: ICE gathering STUN step blocked/delayed WebRTC dial cancellation

replace/ice/gather.go

@@ -435,7 +435,10 @@ func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*stun.UR
 						return
 					}
 
-					xorAddr, err := a.udpMuxSrflx.GetXORMappedAddr(serverAddr, stunGatherTimeout)
+					// [Psiphon]
+					// Pass in context so STUN doesn't block cancelling.
+
+					xorAddr, err := a.udpMuxSrflx.GetXORMappedAddr(ctx, serverAddr, stunGatherTimeout)
 					if err != nil {
 						a.log.Warnf("Failed get server reflexive address %s %s: %v", network, url, err)
 						return

replace/ice/gather_test.go

@@ -841,7 +841,7 @@ func (m *universalUDPMuxMock) GetConnForURL(string, string, net.Addr) (net.Packe
 	return m.conn, nil
 }
 
-func (m *universalUDPMuxMock) GetXORMappedAddr(net.Addr, time.Duration) (*stun.XORMappedAddress, error) {
+func (m *universalUDPMuxMock) GetXORMappedAddr(context.Context, net.Addr, time.Duration) (*stun.XORMappedAddress, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 	m.getXORMappedAddrUsedTimes++

replace/ice/udp_mux_universal.go

@@ -4,6 +4,7 @@
 package ice
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"time"
@@ -18,7 +19,7 @@ import (
 // Actual connection muxing is happening in the UDPMux.
 type UniversalUDPMux interface {
 	UDPMux
-	GetXORMappedAddr(stunAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error)
+	GetXORMappedAddr(ctx context.Context, 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)
 }
@@ -171,7 +172,20 @@ func (m *UniversalUDPMuxDefault) handleXORMappedResponse(stunAddr *net.UDPAddr,
 // 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) {
+func (m *UniversalUDPMuxDefault) GetXORMappedAddr(ctx context.Context, serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
+
+	// [Psiphon]
+	// Add context support so STUN doesn't block cancelling.
+
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	m.mu.Lock()
 	mappedAddr, ok := m.xorMappedMap[serverAddr.String()]
 	// If we already have a mapping for this STUN server (address already received)
@@ -197,6 +211,12 @@ func (m *UniversalUDPMuxDefault) GetXORMappedAddr(serverAddr net.Addr, deadline
 		return nil, fmt.Errorf("%w: %s", errWriteSTUNMessage, err) //nolint:errorlint
 	}
 
+	// [Psiphon]
+	// Add context support so STUN doesn't block cancelling.
+
+	timer := time.NewTimer(deadline)
+	defer timer.Stop()
+
 	// Block until response was handled by the connWorker routine and XORMappedAddress was updated
 	select {
 	case <-waitAddrReceived:
@@ -208,8 +228,10 @@ func (m *UniversalUDPMuxDefault) GetXORMappedAddr(serverAddr net.Addr, deadline
 			return nil, errNoXorAddrMapping
 		}
 		return mappedAddr.addr, nil
-	case <-time.After(deadline):
+	case <-timer.C:
 		return nil, errXORMappedAddrTimeout
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	}
 }
 

replace/ice/udp_mux_universal_test.go

@@ -7,6 +7,7 @@
 package ice
 
 import (
+	"context"
 	"net"
 	"sync"
 	"testing"
@@ -66,7 +67,7 @@ func testMuxSrflxConnection(t *testing.T, udpMux *UniversalUDPMuxDefault, ufrag
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		address, e := udpMux.GetXORMappedAddr(remoteConn.LocalAddr(), time.Second)
+		address, e := udpMux.GetXORMappedAddr(context.Background(), remoteConn.LocalAddr(), time.Second)
 		require.NoError(t, e)
 		require.NotNil(t, address)
 		require.True(t, address.IP.Equal(testXORIP))
@@ -109,7 +110,7 @@ func testMuxSrflxConnection(t *testing.T, udpMux *UniversalUDPMuxDefault, ufrag
 	wg.Wait()
 
 	// We should get address immediately from the cached map
-	address, err := udpMux.GetXORMappedAddr(remoteConn.LocalAddr(), time.Second)
+	address, err := udpMux.GetXORMappedAddr(context.Background(), remoteConn.LocalAddr(), time.Second)
 	require.NoError(t, err)
 	require.NotNil(t, address)
 
@@ -124,7 +125,7 @@ func testMuxSrflxConnection(t *testing.T, udpMux *UniversalUDPMuxDefault, ufrag
 
 	// 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)
+	address, err = udpMux.GetXORMappedAddr(context.Background(), remoteConn.LocalAddr(), time.Millisecond*5)
 	require.NotNil(t, err)
 	require.Nil(t, address)
 }

vendor/github.com/pion/ice/v2/gather.go

@@ -435,7 +435,10 @@ func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*stun.UR
 						return
 					}
 
-					xorAddr, err := a.udpMuxSrflx.GetXORMappedAddr(serverAddr, stunGatherTimeout)
+					// [Psiphon]
+					// Pass in context so STUN doesn't block cancelling.
+
+					xorAddr, err := a.udpMuxSrflx.GetXORMappedAddr(ctx, serverAddr, stunGatherTimeout)
 					if err != nil {
 						a.log.Warnf("Failed get server reflexive address %s %s: %v", network, url, err)
 						return

vendor/github.com/pion/ice/v2/udp_mux_universal.go

@@ -4,6 +4,7 @@
 package ice
 
 import (
+	"context"
 	"fmt"
 	"net"
 	"time"
@@ -18,7 +19,7 @@ import (
 // Actual connection muxing is happening in the UDPMux.
 type UniversalUDPMux interface {
 	UDPMux
-	GetXORMappedAddr(stunAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error)
+	GetXORMappedAddr(ctx context.Context, 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)
 }
@@ -171,7 +172,20 @@ func (m *UniversalUDPMuxDefault) handleXORMappedResponse(stunAddr *net.UDPAddr,
 // 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) {
+func (m *UniversalUDPMuxDefault) GetXORMappedAddr(ctx context.Context, serverAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error) {
+
+	// [Psiphon]
+	// Add context support so STUN doesn't block cancelling.
+
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	default:
+	}
+
 	m.mu.Lock()
 	mappedAddr, ok := m.xorMappedMap[serverAddr.String()]
 	// If we already have a mapping for this STUN server (address already received)
@@ -197,6 +211,12 @@ func (m *UniversalUDPMuxDefault) GetXORMappedAddr(serverAddr net.Addr, deadline
 		return nil, fmt.Errorf("%w: %s", errWriteSTUNMessage, err) //nolint:errorlint
 	}
 
+	// [Psiphon]
+	// Add context support so STUN doesn't block cancelling.
+
+	timer := time.NewTimer(deadline)
+	defer timer.Stop()
+
 	// Block until response was handled by the connWorker routine and XORMappedAddress was updated
 	select {
 	case <-waitAddrReceived:
@@ -208,8 +228,10 @@ func (m *UniversalUDPMuxDefault) GetXORMappedAddr(serverAddr net.Addr, deadline
 			return nil, errNoXorAddrMapping
 		}
 		return mappedAddr.addr, nil
-	case <-time.After(deadline):
+	case <-timer.C:
 		return nil, errXORMappedAddrTimeout
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	}
 }