From 10e34090fa3f0a61f2366ce518955d3b9e168f03 Mon Sep 17 00:00:00 2001 From: Zoa Hickenlooper Date: Tue, 13 Oct 2026 05:21:37 -0400 Subject: [PATCH] [client] Ensure onDisconnected is called on connection teardown (#8178) Conn.Close previously called onDisconnected only when evalStatus reported StatusConnected. When a connection drops or closes from a connecting or disconnecting state, the callback was skipped. This left registered listeners such as Rosenpass unaware of the teardown, keeping peer workers active and leaking sockets. Invoke onDisconnected unconditionally on Conn.Close if registered. --- client/internal/peer/conn.go | 2 +- client/internal/peer/conn_test.go | 43 +++++++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index 17823e043..b8a32b0cb 100644 --- a/client/internal/peer/conn.go +++ b/client/internal/peer/conn.go @@ -341,7 +341,7 @@ func (conn *Conn) Close(signalToRemote bool) { conn.Log.Errorf("failed to remove wg endpoint: %v", err) } - if conn.evalStatus() == StatusConnected && conn.onDisconnected != nil { + if conn.onDisconnected != nil { conn.onDisconnected(conn.config.WgConfig.RemoteKey) } diff --git a/client/internal/peer/conn_test.go b/client/internal/peer/conn_test.go index 8f3d6216d..b481c95ed 100644 --- a/client/internal/peer/conn_test.go +++ b/client/internal/peer/conn_test.go @@ -3,19 +3,26 @@ package peer import ( "context" "fmt" + "net" + "net/netip" "os" "testing" "time" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" "github.com/netbirdio/netbird/client/iface" + "github.com/netbirdio/netbird/client/iface/configurer" + "github.com/netbirdio/netbird/client/iface/wgaddr" + "github.com/netbirdio/netbird/client/iface/wgproxy" "github.com/netbirdio/netbird/client/internal/metrics" "github.com/netbirdio/netbird/client/internal/peer/conntype" "github.com/netbirdio/netbird/client/internal/peer/dispatcher" "github.com/netbirdio/netbird/client/internal/peer/guard" "github.com/netbirdio/netbird/client/internal/peer/ice" + "github.com/netbirdio/netbird/client/internal/peer/worker" "github.com/netbirdio/netbird/client/internal/stdnet" "github.com/netbirdio/netbird/util" ) @@ -418,3 +425,39 @@ func TestMetricsConnType_RelayedMatchesIsRelayed(t *testing.T) { "priority %s: isRelayed and the %q metric tag must agree", priority, tag) } } + +type stubWGIface struct{} + +func (s *stubWGIface) UpdatePeer(peerKey string, allowedIps []netip.Prefix, keepAlive time.Duration, endpoint *net.UDPAddr, preSharedKey *wgtypes.Key) error { + return nil +} +func (s *stubWGIface) RemovePeer(peerKey string) error { return nil } +func (s *stubWGIface) GetStats() (map[string]configurer.WGStats, error) { + return map[string]configurer.WGStats{}, nil +} +func (s *stubWGIface) GetProxy() wgproxy.Proxy { return nil } +func (s *stubWGIface) Address() wgaddr.Address { return wgaddr.Address{} } +func (s *stubWGIface) RemoveEndpointAddress(key string) error { return nil } + +func TestConn_Close_CallsOnDisconnected(t *testing.T) { + var disconnected []string + conn := newWGTimeoutTestConn(true, &disconnected) + conn.endpointUpdater = NewEndpointUpdater(conn.Log, conn.config.WgConfig, false) + conn.config.WgConfig.WgInterface = &stubWGIface{} + conn.endpointUpdater.wgConfig.WgInterface = &stubWGIface{} + conn.workerRelay = &WorkerRelay{} + ctx, cancel := context.WithCancel(context.Background()) + conn.ctx = ctx + conn.ctxCancel = cancel + conn.opened = true + conn.statusICE = worker.NewAtomicStatus() + conn.statusRelay = worker.NewAtomicStatus() + conn.statusRecorder = NewRecorder("https://mgm") + + // The connection is closed while in a disconnected or connecting state + // (e.g. peer dropped, ICE failed, or management disconnect). + conn.Close(false) + + assert.Equal(t, []string{conn.config.WgConfig.RemoteKey}, disconnected, + "Close must notify onDisconnected so Rosenpass cleans up its peer worker") +}