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) + } }