From 210de91665baaae5214a0eaf26d38ae0bc3b58a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zolt=C3=A1n=20Papp?= Date: Thu, 10 Sep 2026 01:02:33 +0200 Subject: [PATCH] [relay] Signal relay disconnects through the conn context AddCloseListener deduplicated listeners by comparing reflect.ValueOf(callback).Pointer(). For a method value that pointer is the address of the compiler-generated wrapper, not an identity bound to the receiver, so every peer's w.onRelayClientDisconnected compared equal. All peers on the home relay register under the same connectionURL key, so only the first registration survived and the rest were silently dropped. On a relay disconnect those peers were never notified: statusRelay stayed connected and the reconnect guard never fired. The relayed net.Conn itself was closed by closeAllConns, so nothing leaked, but the peer state machine did not learn about it. Foreign relays had the same defect scoped to the peers sharing that server. Rather than fixing the deduplication, drop the peer-level listener registry entirely. A relayed Conn now exposes Context(), cancelled when the connection is torn down, with a cancellation cause naming the reason. This is the same shape quic-go uses for its Conn and Stream types, and it removes the whole class of problems around listener identity, lifetime and deregistration: the signal belongs to the resource instead of a side table. WorkerRelay watches that context in a goroutine whose lifetime matches the connection. A watcher that wakes up for a superseded connection compares the conn pointer against the current one and returns without touching the state machine, so a fast relay reconnect cannot have a stale watcher tear down the connection that replaced it. Client.SetOnDisconnectListener stays: it is server-level and drives the reconnect guard and foreign relay eviction, unrelated to peers. handleRelayReady also checks the conn context, closing the race where the relay dies between OpenConn and the readiness handoff and the peer would otherwise build a WireGuard endpoint over a dead connection. TestNotifierDoubleAdd covered the removed mechanism and is gone. TestForeignAutoClose asserted nothing (both branches logged); it now waits for the relay to leave the client map and fails if it does not. --- client/internal/peer/conn.go | 2 +- client/internal/peer/worker_relay.go | 35 +++--- shared/relay/client/client.go | 27 +++-- shared/relay/client/conn.go | 10 ++ shared/relay/client/manager.go | 82 ++------------ shared/relay/client/manager_test.go | 156 ++++++++++++++++++--------- 6 files changed, 163 insertions(+), 149 deletions(-) diff --git a/client/internal/peer/conn.go b/client/internal/peer/conn.go index 83089606f..1ab85adcd 100644 --- a/client/internal/peer/conn.go +++ b/client/internal/peer/conn.go @@ -560,7 +560,7 @@ func (conn *Conn) onRelayConnectionIsReady(rci RelayConnInfo) { conn.mu.Lock() defer conn.mu.Unlock() - if conn.ctx.Err() != nil { + if conn.ctx.Err() != nil || rci.relayedConn.Context().Err() != nil { if err := rci.relayedConn.Close(); err != nil { conn.Log.Warnf("failed to close unnecessary relayed connection: %v", err) } diff --git a/client/internal/peer/worker_relay.go b/client/internal/peer/worker_relay.go index 0402992c9..40d18805b 100644 --- a/client/internal/peer/worker_relay.go +++ b/client/internal/peer/worker_relay.go @@ -3,7 +3,6 @@ package peer import ( "context" "errors" - "net" "net/netip" "sync" "sync/atomic" @@ -14,7 +13,7 @@ import ( ) type RelayConnInfo struct { - relayedConn net.Conn + relayedConn *relayClient.Conn rosenpassPubKey []byte rosenpassAddr string } @@ -27,7 +26,7 @@ type WorkerRelay struct { conn *Conn relayManager *relayClient.Manager - relayedConn net.Conn + relayedConn *relayClient.Conn relayLock sync.Mutex relaySupportedOnRemotePeer atomic.Bool @@ -80,12 +79,7 @@ func (w *WorkerRelay) OnNewOffer(remoteOfferAnswer *OfferAnswer) { w.relayedConn = relayedConn w.relayLock.Unlock() - err = w.relayManager.AddCloseListener(srv, w.onRelayClientDisconnected) - if err != nil { - log.Errorf("failed to add close listener: %s", err) - _ = relayedConn.Close() - return - } + go w.watchRelayedConn(relayedConn) w.log.Debugf("peer conn opened via Relay: %s", srv) go w.conn.onRelayConnectionIsReady(RelayConnInfo{ @@ -109,12 +103,15 @@ func (w *WorkerRelay) RelayIsSupportedLocally() bool { func (w *WorkerRelay) CloseConn() { w.relayLock.Lock() - defer w.relayLock.Unlock() - if w.relayedConn == nil { + conn := w.relayedConn + w.relayedConn = nil + w.relayLock.Unlock() + + if conn == nil { return } - if err := w.relayedConn.Close(); err != nil { + if err := conn.Close(); err != nil { w.log.Warnf("failed to close relay connection: %v", err) } } @@ -133,6 +130,16 @@ func (w *WorkerRelay) preferredRelayServer(myRelayAddress, remoteRelayAddress st return remoteRelayAddress } -func (w *WorkerRelay) onRelayClientDisconnected() { - go w.conn.onRelayDisconnected() +func (w *WorkerRelay) watchRelayedConn(relayedConn *relayClient.Conn) { + <-relayedConn.Context().Done() + + w.relayLock.Lock() + current := w.relayedConn == relayedConn + w.relayLock.Unlock() + + if !current { + return + } + + w.conn.onRelayDisconnected() } diff --git a/shared/relay/client/client.go b/shared/relay/client/client.go index 38c9c7375..1f280db50 100644 --- a/shared/relay/client/client.go +++ b/shared/relay/client/client.go @@ -30,6 +30,12 @@ const ( var ( ErrConnAlreadyExists = fmt.Errorf("connection already exists") + // ErrServerDisconnected is the cancellation cause of a relayed Conn when the + // client lost the connection to the relay server. + ErrServerDisconnected = fmt.Errorf("relay server disconnected") + // ErrPeerDisconnected is the cancellation cause of a relayed Conn when the + // remote peer went offline. + ErrPeerDisconnected = fmt.Errorf("remote peer disconnected") ) type internalStopFlag struct { @@ -74,16 +80,17 @@ type connContainer struct { msgChanLock sync.Mutex closed bool // flag to check if channel is closed ctx context.Context - cancel context.CancelFunc + cancel context.CancelCauseFunc } func newConnContainer(log *log.Entry, c *Client, peerID messages.PeerID, instanceURL *RelayAddr) *connContainer { - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancelCause(context.Background()) msgChan := make(chan Msg, connChannelSize) cn := &Conn{ dstID: peerID, messageChan: msgChan, instanceURL: instanceURL, + ctx: ctx, } cc := &connContainer{ log: log, @@ -128,8 +135,8 @@ func (cc *connContainer) writeMsg(msg Msg) { } } -func (cc *connContainer) close() { - cc.cancel() +func (cc *connContainer) close(cause error) { + cc.cancel(cause) cc.msgChanLock.Lock() defer cc.msgChanLock.Unlock() @@ -298,7 +305,7 @@ func (c *Client) Connect(ctx context.Context) error { // it will return immediately. // It block until the server confirm the peer is online. // todo: what should happen if call with the same peerID with multiple times? -func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, error) { +func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (*Conn, error) { peerID := messages.HashID(dstPeerID) c.mu.Lock() @@ -335,7 +342,7 @@ func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, erro delete(c.conns, peerID) } c.mu.Unlock() - container.close() + container.close(err) return nil, err } @@ -345,7 +352,7 @@ func (c *Client) OpenConn(ctx context.Context, dstPeerID string) (net.Conn, erro delete(c.conns, peerID) } c.mu.Unlock() - container.close() + container.close(ErrServerDisconnected) return nil, fmt.Errorf("relay connection is not established") } c.mu.Unlock() @@ -779,7 +786,7 @@ func (c *Client) listenForStopEvents(ctx context.Context, hc *healthcheck.Receiv func (c *Client) closeAllConns() { for _, container := range c.conns { - container.close() + container.close(ErrServerDisconnected) } c.conns = make(map[messages.PeerID]*connContainer) @@ -799,7 +806,7 @@ func (c *Client) closeConnsByPeerID(peerIDs []messages.PeerID) { } container.log.Infof("remote peer has been disconnected, free up connection: %s", peerID) - container.close() + container.close(ErrPeerDisconnected) delete(c.conns, peerID) } @@ -827,7 +834,7 @@ func (c *Client) closeConn(containerRef *connContainer, id messages.PeerID) erro c.log.Infof("free up connection to peer: %s", id) delete(c.conns, id) - current.close() + current.close(net.ErrClosed) return nil } diff --git a/shared/relay/client/conn.go b/shared/relay/client/conn.go index 9e2279790..67767a2b9 100644 --- a/shared/relay/client/conn.go +++ b/shared/relay/client/conn.go @@ -1,6 +1,7 @@ package client import ( + "context" "net" "time" @@ -12,11 +13,20 @@ type Conn struct { dstID messages.PeerID messageChan chan Msg instanceURL *RelayAddr + ctx context.Context writeFn func(messages.PeerID, []byte) (int, error) closeFn func(messages.PeerID) error localAddrFn func() net.Addr } +// Context returns a context that is cancelled when the connection is torn down, +// either by Close or by the relay client losing the server connection. The +// cancellation cause carries the reason, see ErrServerDisconnected and +// ErrPeerDisconnected. +func (c *Conn) Context() context.Context { + return c.ctx +} + func (c *Conn) Write(p []byte) (n int, err error) { return c.writeFn(c.dstID, p) } diff --git a/shared/relay/client/manager.go b/shared/relay/client/manager.go index 50fcc0b8f..53305d20e 100644 --- a/shared/relay/client/manager.go +++ b/shared/relay/client/manager.go @@ -1,12 +1,9 @@ package client import ( - "container/list" "context" "fmt" - "net" "net/netip" - "reflect" "sync" "time" @@ -43,8 +40,6 @@ func NewRelayTrack() *RelayTrack { } } -type OnServerCloseListener func() - // ManagerOption configures a Manager at construction time. type ManagerOption func(*Manager) @@ -91,7 +86,6 @@ type Manager struct { relayClients map[string]*RelayTrack relayClientsMutex sync.RWMutex - onDisconnectedListeners map[string]*list.List onReconnectedListenerFn func() listenerLock sync.Mutex @@ -126,10 +120,9 @@ func NewManager(ctx context.Context, serverURLs []string, peerID string, mtu uin ConnectionTimeout: defaultConnectionTimeout, TransportFallback: tf, }, - relayClients: make(map[string]*RelayTrack), - onDisconnectedListeners: make(map[string]*list.List), - cleanupInterval: relayCleanupInterval, - keepUnusedServerTime: keepUnusedServerTime, + relayClients: make(map[string]*RelayTrack), + cleanupInterval: relayCleanupInterval, + keepUnusedServerTime: keepUnusedServerTime, } for _, opt := range opts { opt(m) @@ -168,11 +161,11 @@ func (m *Manager) Serve() error { // OpenConn opens a connection to the given peer key. If the peer is on the same relay server, the connection will be // established via the relay server. If the peer is on a different relay server, the manager will establish a new -// connection to the relay server. It returns back with a net.Conn what represent the remote peer connection. +// connection to the relay server. It returns the relayed connection to the remote peer. // // serverIP, when valid and serverAddress is foreign, is used as a dial target if the FQDN-based dial fails. // Ignored for the local home-server path. TLS verification still uses the FQDN via SNI. -func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) { +func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (*Conn, error) { m.relayClientMu.RLock() defer m.relayClientMu.RUnlock() @@ -185,9 +178,7 @@ func (m *Manager) OpenConn(ctx context.Context, serverAddress, peerKey string, s return nil, err } - var ( - netConn net.Conn - ) + var netConn *Conn if !foreign { log.Debugf("open peer connection via permanent server: %s", peerKey) netConn, err = m.relayClient.OpenConn(ctx, peerKey) @@ -220,31 +211,6 @@ func (m *Manager) SetOnReconnectedListener(f func()) { m.onReconnectedListenerFn = f } -// AddCloseListener adds a listener to the given server instance address. The listener will be called if the connection -// closed. -func (m *Manager) AddCloseListener(serverAddress string, onClosedListener OnServerCloseListener) error { - m.relayClientMu.RLock() - defer m.relayClientMu.RUnlock() - - if m.relayClient == nil { - return ErrRelayClientNotConnected - } - - foreign, err := m.isForeignServer(serverAddress) - if err != nil { - return err - } - - var listenerAddr string - if foreign { - listenerAddr = serverAddress - } else { - listenerAddr = m.relayClient.connectionURL - } - m.addListener(listenerAddr, onClosedListener) - return nil -} - // RelayInstanceAddress returns the address and resolved IP of the permanent relay server. It could change if the // network connection is lost. The address is sent to the target peer to choose the common relay server for the // communication; the IP is sent alongside so remote peers can dial directly without their own DNS lookup. Both @@ -334,7 +300,7 @@ func (m *Manager) UpdateToken(token *relayAuth.Token) error { return m.tokenStore.UpdateToken(token) } -func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (net.Conn, error) { +func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string, serverIP netip.Addr) (*Conn, error) { // check if already has a connection to the desired relay server m.relayClientsMutex.RLock() rt, ok := m.relayClients[serverAddress] @@ -387,7 +353,7 @@ func (m *Manager) openConnVia(ctx context.Context, serverAddress, peerKey string // waiting for the dial started by another openConnVia call to finish. It waits // on rt.ready rather than the track lock, so it neither holds nor contends the // track lock across the dial. -func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (net.Conn, error) { +func (m *Manager) openConnOnTrack(ctx context.Context, rt *RelayTrack, peerKey string) (*Conn, error) { select { case <-rt.ready: case <-ctx.Done(): @@ -432,8 +398,6 @@ func (m *Manager) onServerDisconnected(serverAddress string) { if !isHome { m.evictForeignRelay(serverAddress) } - - m.notifyOnDisconnectListeners(serverAddress) } func (m *Manager) evictForeignRelay(serverAddress string) { @@ -527,36 +491,6 @@ func (m *Manager) cleanUpUnusedRelays() { } } -func (m *Manager) addListener(serverAddress string, onClosedListener OnServerCloseListener) { - m.listenerLock.Lock() - defer m.listenerLock.Unlock() - l, ok := m.onDisconnectedListeners[serverAddress] - if !ok { - l = list.New() - } - for e := l.Front(); e != nil; e = e.Next() { - if reflect.ValueOf(e.Value).Pointer() == reflect.ValueOf(onClosedListener).Pointer() { - return - } - } - l.PushBack(onClosedListener) - m.onDisconnectedListeners[serverAddress] = l -} - -func (m *Manager) notifyOnDisconnectListeners(serverAddress string) { - m.listenerLock.Lock() - defer m.listenerLock.Unlock() - - l, ok := m.onDisconnectedListeners[serverAddress] - if !ok { - return - } - for e := l.Front(); e != nil; e = e.Next() { - go e.Value.(OnServerCloseListener)() - } - delete(m.onDisconnectedListeners, serverAddress) -} - func relayConnState(c *Client) RelayConnState { addr, err := c.ServerInstanceURL() if err != nil { diff --git a/shared/relay/client/manager_test.go b/shared/relay/client/manager_test.go index 9e964f688..d6b2e226a 100644 --- a/shared/relay/client/manager_test.go +++ b/shared/relay/client/manager_test.go @@ -2,6 +2,7 @@ package client import ( "context" + "errors" "fmt" "net/netip" "testing" @@ -291,35 +292,29 @@ func TestForeignAutoClose(t *testing.T) { t.Fatalf("failed to serve manager: %s", err) } - // Set up a disconnect listener to track when foreign server disconnects foreignServerURL := toURL(srvCfg2)[0] - disconnected := make(chan struct{}) - onDisconnect := func() { - select { - case disconnected <- struct{}{}: - default: - } - } t.Log("open connection to another peer") if _, err = mgr.OpenConn(ctx, foreignServerURL, "anotherpeer", netip.Addr{}); err == nil { t.Fatalf("should have failed to open connection to another peer") } - // Add the disconnect listener after the connection attempt - if err := mgr.AddCloseListener(foreignServerURL, onDisconnect); err != nil { - t.Logf("failed to add close listener (expected if connection failed): %s", err) - } - - // Wait for cleanup to happen timeout := relayCleanupInterval + keepUnusedServerTime + 2*time.Second t.Logf("waiting for relay cleanup: %s", timeout) - - select { - case <-disconnected: - t.Log("foreign relay connection cleaned up successfully") - case <-time.After(timeout): - t.Log("timeout waiting for cleanup - this might be expected if connection never established") + deadline := time.After(timeout) + for { + mgr.relayClientsMutex.RLock() + _, tracked := mgr.relayClients[foreignServerURL] + mgr.relayClientsMutex.RUnlock() + if !tracked { + t.Log("foreign relay connection cleaned up successfully") + break + } + select { + case <-deadline: + t.Fatal("foreign relay was not cleaned up") + case <-time.After(200 * time.Millisecond): + } } t.Logf("closing manager") @@ -413,23 +408,24 @@ func waitForReady(ctx context.Context, m *Manager, timeout time.Duration) error return fmt.Errorf("manager not ready within %s", timeout) } -func TestNotifierDoubleAdd(t *testing.T) { +func toURL(address server.ListenerConfig) []string { + return []string{"rel://" + address.Address} +} + +func TestConnContextCancelledOnServerDisconnect(t *testing.T) { ctx := context.Background() - listenerCfg1 := server.ListenerConfig{ - Address: "localhost:52501", - } - srv, err := server.NewServer(newManagerTestServerConfig(listenerCfg1.Address)) + srvCfg := server.ListenerConfig{Address: "localhost:52601"} + srv, err := server.NewServer(newManagerTestServerConfig(srvCfg.Address)) if err != nil { t.Fatalf("failed to create server: %s", err) } errChan := make(chan error, 1) go func() { - if err := srv.Listen(listenerCfg1); err != nil { + if err := srv.Listen(srvCfg); err != nil { errChan <- err } }() - defer func() { if err := srv.Shutdown(ctx); err != nil { t.Errorf("failed to close server: %s", err) @@ -440,46 +436,106 @@ func TestNotifierDoubleAdd(t *testing.T) { t.Fatalf("failed to start server: %s", err) } - log.Debugf("connect by alice") mCtx, cancel := context.WithCancel(ctx) defer cancel() - clientBob := NewManager(mCtx, toURL(listenerCfg1), "bob", iface.DefaultMTU) - if err = clientBob.Serve(); err != nil { + mgrBob := NewManager(mCtx, toURL(srvCfg), "bob", iface.DefaultMTU) + if err := mgrBob.Serve(); err != nil { + t.Fatalf("failed to serve bob manager: %s", err) + } + + mgr := NewManager(mCtx, toURL(srvCfg), "alice", iface.DefaultMTU) + if err := mgr.Serve(); err != nil { t.Fatalf("failed to serve manager: %s", err) } - clientAlice := NewManager(mCtx, toURL(listenerCfg1), "alice", iface.DefaultMTU) - if err = clientAlice.Serve(); err != nil { + ra, _, err := mgr.RelayInstanceAddress() + if err != nil { + t.Fatalf("failed to get relay address: %s", err) + } + + relayedConn, err := mgr.OpenConn(ctx, ra, "bob", netip.Addr{}) + if err != nil { + t.Fatalf("failed to open conn: %s", err) + } + + select { + case <-relayedConn.Context().Done(): + t.Fatal("conn context cancelled while the relay is still up") + default: + } + + _ = mgr.relayClient.relayConn.Close() + + select { + case <-relayedConn.Context().Done(): + case <-time.After(15 * time.Second): + t.Fatal("conn context was not cancelled after the relay connection dropped") + } + + if cause := context.Cause(relayedConn.Context()); !errors.Is(cause, ErrServerDisconnected) { + t.Errorf("unexpected cancellation cause: %v, want %v", cause, ErrServerDisconnected) + } +} + +func TestConnContextCauseOnLocalClose(t *testing.T) { + ctx := context.Background() + + srvCfg := server.ListenerConfig{Address: "localhost:52602"} + srv, err := server.NewServer(newManagerTestServerConfig(srvCfg.Address)) + if err != nil { + t.Fatalf("failed to create server: %s", err) + } + errChan := make(chan error, 1) + go func() { + if err := srv.Listen(srvCfg); err != nil { + errChan <- err + } + }() + defer func() { + if err := srv.Shutdown(ctx); err != nil { + t.Errorf("failed to close server: %s", err) + } + }() + + if err := waitForServerToStart(errChan); err != nil { + t.Fatalf("failed to start server: %s", err) + } + + mCtx, cancel := context.WithCancel(ctx) + defer cancel() + + mgrBob := NewManager(mCtx, toURL(srvCfg), "bob", iface.DefaultMTU) + if err := mgrBob.Serve(); err != nil { + t.Fatalf("failed to serve bob manager: %s", err) + } + + mgr := NewManager(mCtx, toURL(srvCfg), "alice", iface.DefaultMTU) + if err := mgr.Serve(); err != nil { t.Fatalf("failed to serve manager: %s", err) } - conn1, err := clientAlice.OpenConn(ctx, clientAlice.ServerURLs()[0], "bob", netip.Addr{}) + ra, _, err := mgr.RelayInstanceAddress() if err != nil { - t.Fatalf("failed to bind channel: %s", err) + t.Fatalf("failed to get relay address: %s", err) } - fnCloseListener := OnServerCloseListener(func() { - log.Infof("close listener") - }) - - err = clientAlice.AddCloseListener(clientAlice.ServerURLs()[0], fnCloseListener) + relayedConn, err := mgr.OpenConn(ctx, ra, "bob", netip.Addr{}) if err != nil { - t.Fatalf("failed to add close listener: %s", err) + t.Fatalf("failed to open conn: %s", err) } - err = clientAlice.AddCloseListener(clientAlice.ServerURLs()[0], fnCloseListener) - if err != nil { - t.Fatalf("failed to add close listener: %s", err) + if err := relayedConn.Close(); err != nil { + t.Fatalf("failed to close conn: %s", err) } - err = conn1.Close() - if err != nil { - t.Errorf("failed to close connection: %s", err) + select { + case <-relayedConn.Context().Done(): + case <-time.After(5 * time.Second): + t.Fatal("conn context was not cancelled after a local close") } -} - -func toURL(address server.ListenerConfig) []string { - return []string{"rel://" + address.Address} + if cause := context.Cause(relayedConn.Context()); errors.Is(cause, ErrServerDisconnected) { + t.Errorf("local close must not be reported as a server disconnect, got: %v", cause) + } }