diff --git a/client/iface/bind/ice_bind.go b/client/iface/bind/ice_bind.go index 41f415af7..c203b5bfc 100644 --- a/client/iface/bind/ice_bind.go +++ b/client/iface/bind/ice_bind.go @@ -1,6 +1,7 @@ package bind import ( + "context" "fmt" "net" "net/netip" @@ -38,7 +39,7 @@ func (rc receiverCreator) CreateIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UD // use the port because in the Send function the wgConn.Endpoint the port info is not exported. type ICEBind struct { *wgConn.StdNetBind - RecvChan chan RecvMessage + recvChan chan RecvMessage transportNet transport.Net filterFn FilterFn @@ -58,7 +59,7 @@ func NewICEBind(transportNet transport.Net, filterFn FilterFn) *ICEBind { b, _ := wgConn.NewStdNetBind().(*wgConn.StdNetBind) ib := &ICEBind{ StdNetBind: b, - RecvChan: make(chan RecvMessage, 1), + recvChan: make(chan RecvMessage, 1), transportNet: transportNet, filterFn: filterFn, endpoints: make(map[netip.Addr]net.Conn), @@ -155,6 +156,14 @@ func (b *ICEBind) Send(bufs [][]byte, ep wgConn.Endpoint) error { return nil } +func (b *ICEBind) Recv(ctx context.Context, msg RecvMessage) { + select { + case <-ctx.Done(): + return + case b.recvChan <- msg: + } +} + func (s *ICEBind) createIPv4ReceiverFn(pc *ipv4.PacketConn, conn *net.UDPConn, rxOffload bool, msgsPool *sync.Pool) wgConn.ReceiveFunc { s.muUDPMux.Lock() defer s.muUDPMux.Unlock() @@ -264,7 +273,7 @@ func (c *ICEBind) receiveRelayed(buffs [][]byte, sizes []int, eps []wgConn.Endpo select { case <-c.closedChan: return 0, net.ErrClosed - case msg, ok := <-c.RecvChan: + case msg, ok := <-c.recvChan: if !ok { return 0, net.ErrClosed } diff --git a/client/iface/wgproxy/bind/proxy.go b/client/iface/wgproxy/bind/proxy.go index c1cc8dd30..b0dfef06b 100644 --- a/client/iface/wgproxy/bind/proxy.go +++ b/client/iface/wgproxy/bind/proxy.go @@ -13,8 +13,14 @@ import ( "github.com/netbirdio/netbird/client/iface/bind" ) +type IceBind interface { + SetEndpoint(addr *net.UDPAddr, conn net.Conn) (*net.UDPAddr, error) + RemoveEndpoint(addr *net.UDPAddr) + Recv(ctx context.Context, msg bind.RecvMessage) +} + type ProxyBind struct { - bind *bind.ICEBind + bind IceBind // wgEndpoint is a fake address that generated by the Bind.SetEndpoint based on the remote NetBird peer address wgRelayedEndpoint *bind.Endpoint @@ -30,7 +36,7 @@ type ProxyBind struct { isStarted bool } -func NewProxyBind(bind *bind.ICEBind) *ProxyBind { +func NewProxyBind(bind IceBind) *ProxyBind { return &ProxyBind{ bind: bind, pausedCond: sync.NewCond(&sync.Mutex{}), @@ -172,7 +178,7 @@ func (p *ProxyBind) proxyToLocal(ctx context.Context) { Endpoint: p.wgCurrentUsed, Buffer: buf[:n], } - p.bind.RecvChan <- msg + p.bind.Recv(ctx, msg) p.pausedCond.L.Unlock() } }