From fc156259630ad0a9e81bebb83c17713a5d45acd5 Mon Sep 17 00:00:00 2001
From: Viktor Liu <viktor@netbird.io>
Date: Thu, 23 Oct 2025 15:06:53 +0200
Subject: [PATCH] Clean up failed conn hooks

---
 client/net/conn.go               | 12 +++++++-----
 client/net/dial.go               |  1 -
 client/net/dialer_dial.go        |  3 ++-
 client/net/listener_listen.go    |  4 ++--
 shared/management/client/grpc.go |  3 +--
 shared/signal/client/grpc.go     |  3 +--
 6 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/client/net/conn.go b/client/net/conn.go
index 918e7f628..bf54c792d 100644
--- a/client/net/conn.go
+++ b/client/net/conn.go
@@ -17,8 +17,7 @@ type Conn struct {
 	ID hooks.ConnectionID
 }
 
-// Close overrides the net.Conn Close method to execute all registered hooks after closing the connection
-// Close overrides the net.Conn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.Conn Close method to execute all registered hooks after closing the connection.
 func (c *Conn) Close() error {
 	return closeConn(c.ID, c.Conn)
 }
@@ -29,7 +28,7 @@ type TCPConn struct {
 	ID hooks.ConnectionID
 }
 
-// Close overrides the net.TCPConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.TCPConn Close method to execute all registered hooks after closing the connection.
 func (c *TCPConn) Close() error {
 	return closeConn(c.ID, c.TCPConn)
 }
@@ -37,13 +36,16 @@ func (c *TCPConn) Close() error {
 // closeConn is a helper function to close connections and execute close hooks.
 func closeConn(id hooks.ConnectionID, conn io.Closer) error {
 	err := conn.Close()
+	cleanupConnID(id)
+	return err
+}
 
+// cleanupConnID executes close hooks for a connection ID.
+func cleanupConnID(id hooks.ConnectionID) {
 	closeHooks := hooks.GetCloseHooks()
 	for _, hook := range closeHooks {
 		if err := hook(id); err != nil {
 			log.Errorf("Error executing close hook: %v", err)
 		}
 	}
-
-	return err
 }
diff --git a/client/net/dial.go b/client/net/dial.go
index 041a00e5d..17c9ff98a 100644
--- a/client/net/dial.go
+++ b/client/net/dial.go
@@ -74,7 +74,6 @@ func DialTCP(network string, laddr, raddr *net.TCPAddr) (transport.TCPConn, erro
 		}
 		return &TCPConn{TCPConn: tcpConn, ID: c.ID}, nil
 	}
-
 	if err := conn.Close(); err != nil {
 		log.Errorf("failed to close connection: %v", err)
 	}
diff --git a/client/net/dialer_dial.go b/client/net/dialer_dial.go
index 2e1eb53d8..1e275013f 100644
--- a/client/net/dialer_dial.go
+++ b/client/net/dialer_dial.go
@@ -30,6 +30,7 @@ func (d *Dialer) DialContext(ctx context.Context, network, address string) (net.
 
 	conn, err := d.Dialer.DialContext(ctx, network, address)
 	if err != nil {
+		cleanupConnID(connID)
 		return nil, fmt.Errorf("d.Dialer.DialContext: %w", err)
 	}
 
@@ -64,7 +65,7 @@ func callDialerHooks(ctx context.Context, connID hooks.ConnectionID, address str
 
 	ips, err := resolver.LookupIPAddr(ctx, host)
 	if err != nil {
-		return fmt.Errorf("failed to resolve address %s: %w", address, err)
+		return fmt.Errorf("resolve address %s: %w", address, err)
 	}
 
 	log.Debugf("Dialer resolved IPs for %s: %v", address, ips)
diff --git a/client/net/listener_listen.go b/client/net/listener_listen.go
index 0bb5ad67d..a150172b4 100644
--- a/client/net/listener_listen.go
+++ b/client/net/listener_listen.go
@@ -48,7 +48,7 @@ func (c *PacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return c.PacketConn.WriteTo(b, addr)
 }
 
-// Close overrides the net.PacketConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.PacketConn Close method to execute all registered hooks after closing the connection.
 func (c *PacketConn) Close() error {
 	defer c.seenAddrs.Clear()
 	return closeConn(c.ID, c.PacketConn)
@@ -69,7 +69,7 @@ func (c *UDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	return c.UDPConn.WriteTo(b, addr)
 }
 
-// Close overrides the net.UDPConn Close method to execute all registered hooks before closing the connection.
+// Close overrides the net.UDPConn Close method to execute all registered hooks after closing the connection.
 func (c *UDPConn) Close() error {
 	defer c.seenAddrs.Clear()
 	return closeConn(c.ID, c.UDPConn)
diff --git a/shared/management/client/grpc.go b/shared/management/client/grpc.go
index 076f2532b..520a83e36 100644
--- a/shared/management/client/grpc.go
+++ b/shared/management/client/grpc.go
@@ -55,8 +55,7 @@ func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsE
 		var err error
 		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.ManagementComponent)
 		if err != nil {
-			log.Printf("createConnection error: %v", err)
-			return err
+			return fmt.Errorf("create connection: %w", err)
 		}
 		return nil
 	}
diff --git a/shared/signal/client/grpc.go b/shared/signal/client/grpc.go
index 31f3372c0..5368b57a2 100644
--- a/shared/signal/client/grpc.go
+++ b/shared/signal/client/grpc.go
@@ -60,8 +60,7 @@ func NewClient(ctx context.Context, addr string, key wgtypes.Key, tlsEnabled boo
 		var err error
 		conn, err = nbgrpc.CreateConnection(ctx, addr, tlsEnabled, wsproxy.SignalComponent)
 		if err != nil {
-			log.Printf("createConnection error: %v", err)
-			return err
+			return fmt.Errorf("create connection: %w", err)
 		}
 		return nil
 	}