From ffe74365a83cd922797296791fe091c6c3adad18 Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Fri, 14 Feb 2025 15:54:40 +0100
Subject: [PATCH] Code cleaning

---
 client/iface/bind/endpoint.go      | 13 ++++++++++++-
 client/iface/wgproxy/bind/proxy.go | 18 +++++++++++-------
 2 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/client/iface/bind/endpoint.go b/client/iface/bind/endpoint.go
index 1926ff88f..bce2460de 100644
--- a/client/iface/bind/endpoint.go
+++ b/client/iface/bind/endpoint.go
@@ -1,5 +1,16 @@
 package bind
 
-import wgConn "golang.zx2c4.com/wireguard/conn"
+import (
+	"net"
+
+	wgConn "golang.zx2c4.com/wireguard/conn"
+)
 
 type Endpoint = wgConn.StdNetEndpoint
+
+func EndpointToUDPAddr(e Endpoint) *net.UDPAddr {
+	return &net.UDPAddr{
+		IP:   e.Addr().AsSlice(),
+		Port: int(e.Port()),
+	}
+}
diff --git a/client/iface/wgproxy/bind/proxy.go b/client/iface/wgproxy/bind/proxy.go
index 8a2e65382..1d5390a57 100644
--- a/client/iface/wgproxy/bind/proxy.go
+++ b/client/iface/wgproxy/bind/proxy.go
@@ -16,7 +16,7 @@ import (
 type ProxyBind struct {
 	Bind *bind.ICEBind
 
-	wgAddr     *net.UDPAddr
+	// wgEndpoint is a fake address that generated by the Bind.SetEndpoint based on the remote NetBird peer address
 	wgEndpoint *bind.Endpoint
 	remoteConn net.Conn
 	ctx        context.Context
@@ -32,21 +32,25 @@ type ProxyBind struct {
 // AddTurnConn adds a new connection to the bind.
 // endpoint is the NetBird address of the remote peer. The SetEndpoint return with the address what will be used in the
 // WireGuard configuration.
+//
+// Parameters:
+//   - ctx: Context is used for proxyToLocal to avoid unnecessary error messages
+//   - nbAddr: The NetBird UDP address of the remote peer, it required to generate fake address
+//   - remoteConn: The established TURN connection to the remote peer
 func (p *ProxyBind) AddTurnConn(ctx context.Context, nbAddr *net.UDPAddr, remoteConn net.Conn) error {
-	addr, err := p.Bind.SetEndpoint(nbAddr, remoteConn)
+	fakeAddr, err := p.Bind.SetEndpoint(nbAddr, remoteConn)
 	if err != nil {
 		return err
 	}
 
-	p.wgAddr = addr
-	p.wgEndpoint = addrToEndpoint(addr)
+	p.wgEndpoint = addrToEndpoint(fakeAddr)
 	p.remoteConn = remoteConn
 	p.ctx, p.cancel = context.WithCancel(ctx)
 	return err
-
 }
+
 func (p *ProxyBind) EndpointAddr() *net.UDPAddr {
-	return p.wgAddr
+	return bind.EndpointToUDPAddr(*p.wgEndpoint)
 }
 
 func (p *ProxyBind) Work() {
@@ -93,7 +97,7 @@ func (p *ProxyBind) close() error {
 
 	p.cancel()
 
-	p.Bind.RemoveEndpoint(p.wgAddr)
+	p.Bind.RemoveEndpoint(bind.EndpointToUDPAddr(*p.wgEndpoint))
 
 	if rErr := p.remoteConn.Close(); rErr != nil && !errors.Is(rErr, net.ErrClosed) {
 		return rErr