Process routes before peers (#2105)

iface/bind/bind.go

@@ -28,11 +28,14 @@ type ICEBind struct {
 
 	transportNet transport.Net
 	udpMux       *UniversalUDPMuxDefault
+
+	filterFn FilterFn
 }
 
-func NewICEBind(transportNet transport.Net) *ICEBind {
+func NewICEBind(transportNet transport.Net, filterFn FilterFn) *ICEBind {
 	ib := &ICEBind{
 		transportNet: transportNet,
+		filterFn:     filterFn,
 	}
 
 	rc := receiverCreator{
@@ -59,8 +62,9 @@ func (s *ICEBind) createIPv4ReceiverFn(ipv4MsgsPool *sync.Pool, pc *ipv4.PacketC
 
 	s.udpMux = NewUniversalUDPMuxDefault(
 		UniversalUDPMuxParams{
-			UDPConn: conn,
-			Net:     s.transportNet,
+			UDPConn:  conn,
+			Net:      s.transportNet,
+			FilterFn: s.filterFn,
 		},
 	)
 	return func(bufs [][]byte, sizes []int, eps []wgConn.Endpoint) (n int, err error) {

iface/bind/udp_mux_universal.go

@@ -8,6 +8,8 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"net/netip"
+	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -17,6 +19,10 @@ import (
 	"github.com/pion/transport/v3"
 )
 
+// FilterFn is a function that filters out candidates based on the address.
+// If it returns true, the address is to be filtered. It also returns the prefix of matching route.
+type FilterFn func(address netip.Addr) (bool, netip.Prefix, error)
+
 // UniversalUDPMuxDefault handles STUN and TURN servers packets by wrapping the original UDPConn
 // It then passes packets to the UDPMux that does the actual connection muxing.
 type UniversalUDPMuxDefault struct {
@@ -34,6 +40,7 @@ type UniversalUDPMuxParams struct {
 	UDPConn               net.PacketConn
 	XORMappedAddrCacheTTL time.Duration
 	Net                   transport.Net
+	FilterFn              FilterFn
 }
 
 // NewUniversalUDPMuxDefault creates an implementation of UniversalUDPMux embedding UDPMux
@@ -56,6 +63,7 @@ func NewUniversalUDPMuxDefault(params UniversalUDPMuxParams) *UniversalUDPMuxDef
 		PacketConn: params.UDPConn,
 		mux:        m,
 		logger:     params.Logger,
+		filterFn:   params.FilterFn,
 	}
 
 	// embed UDPMux
@@ -105,8 +113,68 @@ func (m *UniversalUDPMuxDefault) ReadFromConn(ctx context.Context) {
 // 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
+	mux      *UniversalUDPMuxDefault
+	logger   logging.LeveledLogger
+	filterFn FilterFn
+	// TODO: reset cache on route changes
+	addrCache sync.Map
+}
+
+func (u *udpConn) WriteTo(b []byte, addr net.Addr) (int, error) {
+	if u.filterFn == nil {
+		return u.PacketConn.WriteTo(b, addr)
+	}
+
+	if isRouted, found := u.addrCache.Load(addr.String()); found {
+		return u.handleCachedAddress(isRouted.(bool), b, addr)
+	}
+
+	return u.handleUncachedAddress(b, addr)
+}
+
+func (u *udpConn) handleCachedAddress(isRouted bool, b []byte, addr net.Addr) (int, error) {
+	if isRouted {
+		return 0, fmt.Errorf("address %s is part of a routed network, refusing to write", addr)
+	}
+	return u.PacketConn.WriteTo(b, addr)
+}
+
+func (u *udpConn) handleUncachedAddress(b []byte, addr net.Addr) (int, error) {
+	if err := u.performFilterCheck(addr); err != nil {
+		return 0, err
+	}
+	return u.PacketConn.WriteTo(b, addr)
+}
+
+func (u *udpConn) performFilterCheck(addr net.Addr) error {
+	host, err := getHostFromAddr(addr)
+	if err != nil {
+		log.Errorf("Failed to get host from address %s: %v", addr, err)
+		return nil
+	}
+
+	a, err := netip.ParseAddr(host)
+	if err != nil {
+		log.Errorf("Failed to parse address %s: %v", addr, err)
+		return nil
+	}
+
+	if isRouted, prefix, err := u.filterFn(a); err != nil {
+		log.Errorf("Failed to check if address %s is routed: %v", addr, err)
+	} else {
+		u.addrCache.Store(addr.String(), isRouted)
+		if isRouted {
+			// Extra log, as the error only shows up with ICE logging enabled
+			log.Infof("Address %s is part of routed network %s, refusing to write", addr, prefix)
+			return fmt.Errorf("address %s is part of routed network %s, refusing to write", addr, prefix)
+		}
+	}
+	return nil
+}
+
+func getHostFromAddr(addr net.Addr) (string, error) {
+	host, _, err := net.SplitHostPort(addr.String())
+	return host, err
 }
 
 // GetSharedConn returns the shared udp conn