diff --git a/libp2p-hs.cabal b/libp2p-hs.cabal index 7a7c0d8..e8b6ca9 100644 --- a/libp2p-hs.cabal +++ b/libp2p-hs.cabal @@ -202,6 +202,7 @@ test-suite libp2p-hs-test LibP2P.NAT.Relay.RelaySpec LibP2P.NAT.Relay.ClientSpec LibP2P.NAT.Relay.TransportSpec + LibP2P.NAT.Relay.ReservationLifecycleSpec LibP2P.NAT.DCUtR.MessageSpec LibP2P.NAT.DCUtR.DCUtRSpec LibP2P.NAT.RegistrationSpec diff --git a/src/LibP2P/NAT.hs b/src/LibP2P/NAT.hs index 2021de3..a346542 100644 --- a/src/LibP2P/NAT.hs +++ b/src/LibP2P/NAT.hs @@ -18,11 +18,13 @@ module LibP2P.NAT , registerRelayHopHandler , registerRelayStopHandler , registerDCUtRHandler + , registerReservationCleanup -- * Circuit client , CircuitState ) where -import Control.Concurrent.STM (atomically) +import Control.Concurrent.STM (atomically, modifyTVar') +import qualified Data.Map.Strict as Map import Control.Exception (SomeException, catch, try) import LibP2P.Crypto.PeerId (PeerId, peerIdBytes) import LibP2P.Multiaddr (Multiaddr (..), encapsulate) @@ -44,6 +46,7 @@ import LibP2P.NAT.Relay , handleConnect , handleReserve , newRelayState + , rsReservations ) import LibP2P.NAT.Relay.Client (handleStop) import LibP2P.NAT.Relay.Message @@ -103,8 +106,34 @@ registerNATHandlers sw config = do registerRelayHopHandler sw relayState registerRelayStopHandler sw circuitState registerDCUtRHandler sw + registerReservationCleanup sw relayState pure (relayState, circuitState) +-- | Drop a peer's relay reservation once its last connection to us goes +-- away (specs/relay/circuit-v2): "the reservation remains valid until +-- its expiration, as long as there is an active connection from the peer +-- to the relay. If the peer disconnects, the reservation is no longer +-- valid." +-- +-- The reservation is bound to the peer, not to the connection the +-- RESERVE arrived on, so a peer holding a second connection keeps it. +-- This matches go-libp2p, whose relay returns early from its disconnect +-- notifiee while @Connectedness(p) == Connected@. +-- +-- 'closeConnection' removes the connection from the pool in the same STM +-- transaction that marks it closed, and only then runs the notifiers, so +-- the lookup below never observes the connection being torn down. +registerReservationCleanup :: Switch -> RelayState -> IO () +registerReservationCleanup sw relayState = + atomically $ modifyTVar' (swDisconnectNotifiers sw) (dropReservation :) + where + dropReservation conn = atomically $ do + let peerId = connPeerId conn + remaining <- lookupConn (swConnPool sw) peerId + case remaining of + Just _ -> pure () + Nothing -> modifyTVar' (rsReservations relayState) (Map.delete peerId) + -- | Register the AutoNAT server handler (/libp2p/autonat/1.0.0). -- -- The dial-back deliberately bypasses the connection pool: reusing the diff --git a/src/LibP2P/Switch.hs b/src/LibP2P/Switch.hs index 1d870ed..8444dfb 100644 --- a/src/LibP2P/Switch.hs +++ b/src/LibP2P/Switch.hs @@ -46,6 +46,7 @@ newSwitch pid kp = do } peerStoreVar <- newTVarIO Map.empty notifiersVar <- newTVarIO [] + disconnectNotifiersVar <- newTVarIO [] listenersVar <- newTVarIO [] pure Switch { swLocalPeerId = pid @@ -60,6 +61,7 @@ newSwitch pid kp = do , swResourceMgr = resMgr , swPeerStore = peerStoreVar , swNotifiers = notifiersVar + , swDisconnectNotifiers = disconnectNotifiersVar , swListeners = listenersVar } diff --git a/src/LibP2P/Switch/Connection.hs b/src/LibP2P/Switch/Connection.hs index d651f4d..6960991 100644 --- a/src/LibP2P/Switch/Connection.hs +++ b/src/LibP2P/Switch/Connection.hs @@ -33,12 +33,19 @@ import LibP2P.Switch.Types ) -- | Tear down a connection: remove it from the pool, release its --- resource reservation, publish a Disconnected event, and close the --- muxer session together with the underlying transport. +-- resource reservation, publish a Disconnected event, run the +-- disconnect notifiers, and close the muxer session together with the +-- underlying transport. -- -- Idempotent: the state transition to ConnClosed is atomic, so -- concurrent calls (accept loop exit, explicit close, switchClose) --- perform the teardown exactly once. +-- perform the teardown exactly once, and the notifiers run once. +-- +-- Notifiers run synchronously, after the pool removal has committed and +-- before the muxer is closed, so a notifier that asks whether the peer +-- still has a live connection (Circuit Relay v2 reservation cleanup) +-- never sees the connection being torn down. Each is isolated so a +-- failing notifier cannot abort the teardown. closeConnection :: Switch -> Connection -> IO () closeConnection sw conn = do shouldClose <- atomically $ do @@ -52,7 +59,9 @@ closeConnection sw conn = do writeTChan (swEvents sw) (Disconnected (connPeerId conn) (connDirection conn) (connRemoteAddr conn)) pure True - when shouldClose $ + when shouldClose $ do + notifiers <- atomically $ readTVar (swDisconnectNotifiers sw) + mapM_ (\f -> f conn `catch` \(_ :: SomeException) -> pure ()) notifiers muxClose (connSession conn) `catch` \(_ :: SomeException) -> pure () -- | Tear down every pooled connection (used by switchClose). diff --git a/src/LibP2P/Switch/Types.hs b/src/LibP2P/Switch/Types.hs index 15c09ff..6f5864b 100644 --- a/src/LibP2P/Switch/Types.hs +++ b/src/LibP2P/Switch/Types.hs @@ -120,5 +120,6 @@ data Switch = Switch , swResourceMgr :: !ResourceManager -- ^ Hierarchical resource manager , swPeerStore :: !(TVar (Map PeerId IdentifyInfo)) -- ^ Identify info per peer , swNotifiers :: !(TVar [Connection -> IO ()]) -- ^ Callbacks on new connection + , swDisconnectNotifiers :: !(TVar [Connection -> IO ()]) -- ^ Callbacks on connection teardown , swListeners :: !(TVar [ActiveListener]) -- ^ Active listeners } diff --git a/test/LibP2P/DHT/APISpec.hs b/test/LibP2P/DHT/APISpec.hs index 8e116b7..59f1aa2 100644 --- a/test/LibP2P/DHT/APISpec.hs +++ b/test/LibP2P/DHT/APISpec.hs @@ -198,6 +198,7 @@ mkMockSwitch pid = do resMgr <- mkMockResourceMgr peerStore <- newTVarIO Map.empty notifiers <- newTVarIO [] + disconnectNotifiers <- newTVarIO [] listeners <- newTVarIO [] kp <- getDummyKeyPair pure Switch @@ -213,6 +214,7 @@ mkMockSwitch pid = do , swResourceMgr = resMgr , swPeerStore = peerStore , swNotifiers = notifiers + , swDisconnectNotifiers = disconnectNotifiers , swListeners = listeners } diff --git a/test/LibP2P/DHT/DHTSpec.hs b/test/LibP2P/DHT/DHTSpec.hs index 3d96d73..495eb79 100644 --- a/test/LibP2P/DHT/DHTSpec.hs +++ b/test/LibP2P/DHT/DHTSpec.hs @@ -70,6 +70,7 @@ mkMockSwitch pid = do resMgr <- mkMockResourceMgr peerStore <- newTVarIO Map.empty notifiers <- newTVarIO [] + disconnectNotifiers <- newTVarIO [] listeners <- newTVarIO [] pure Switch { swLocalPeerId = pid @@ -84,6 +85,7 @@ mkMockSwitch pid = do , swResourceMgr = resMgr , swPeerStore = peerStore , swNotifiers = notifiers + , swDisconnectNotifiers = disconnectNotifiers , swListeners = listeners } diff --git a/test/LibP2P/NAT/Relay/ReservationLifecycleSpec.hs b/test/LibP2P/NAT/Relay/ReservationLifecycleSpec.hs new file mode 100644 index 0000000..4280c45 --- /dev/null +++ b/test/LibP2P/NAT/Relay/ReservationLifecycleSpec.hs @@ -0,0 +1,208 @@ +-- | Tests for relay reservation invalidation on disconnect (issue #255). +-- +-- specs/relay/circuit-v2: "the reservation remains valid until its +-- expiration, as long as there is an active connection from the peer to +-- the relay. If the peer disconnects, the reservation is no longer +-- valid." +-- +-- The reservation is bound to the peer, not to the connection the +-- RESERVE arrived on, so these tests pin down both halves: a surviving +-- connection keeps the reservation alive, and losing the last one drops +-- it. Teardown is driven by calling 'closeConnection' on the relay's own +-- side of the connection, which runs the disconnect notifiers +-- synchronously and keeps the assertions deterministic. +module LibP2P.NAT.Relay.ReservationLifecycleSpec (spec) where + +import Control.Concurrent (threadDelay) +import Control.Concurrent.STM (atomically, modifyTVar', readTVarIO) +import qualified Data.Map.Strict as Map +import LibP2P.Crypto.Ed25519 (generateKeyPair) +import LibP2P.Crypto.Key (KeyPair, publicKey) +import LibP2P.Crypto.PeerId (PeerId, fromPublicKey) +import LibP2P.Multiaddr (Multiaddr (..)) +import LibP2P.Multiaddr.Protocol (Protocol (..)) +import LibP2P.MultistreamSelect.Negotiation + ( NegotiationResult (..) + , negotiateInitiator + ) +import LibP2P.NAT (NATConfig (..), defaultNATConfig, registerNATHandlers) +import LibP2P.NAT.Relay + ( ActiveReservation (..) + , RelayConfig (..) + , RelayState (..) + , defaultRelayConfig + ) +import LibP2P.NAT.Relay.Client (makeReservation) +import LibP2P.NAT.Relay.Message (HopMessage (..), RelayStatus (..), hopProtocolId) +import LibP2P.Switch (addTransport, newSwitch, switchClose) +import LibP2P.Switch.ConnPool (lookupAllConns) +import LibP2P.Switch.Connection (closeConnection, newStream) +import LibP2P.Switch.Dial (dial) +import LibP2P.Switch.Listen (defaultConnectionGater, switchListen) +import LibP2P.Switch.Types (Connection (..), Switch (..)) +import LibP2P.Switch.Upgrade (upgradeOutbound) +import LibP2P.Transport (Transport (..)) +import LibP2P.Transport.TCP (newTCPTransport) +import Test.Hspec + +-- | Generate a test identity (PeerId, KeyPair). +mkTestIdentity :: IO (PeerId, KeyPair) +mkTestIdentity = do + Right kp <- generateKeyPair + let pid = fromPublicKey (publicKey kp) + pure (pid, kp) + +-- | Loopback address with port 0 (OS assigns ephemeral port). +loopbackAddr :: Multiaddr +loopbackAddr = Multiaddr [IP4 0x7f000001, TCP 0] + +-- | A switch with TCP and a listener, used for the relay and its clients. +newListeningSwitch :: IO (Switch, PeerId, [Multiaddr]) +newListeningSwitch = do + (pid, kp) <- mkTestIdentity + sw <- newSwitch pid kp + addTransport sw =<< newTCPTransport + addrs <- switchListen sw defaultConnectionGater [loopbackAddr] + pure (sw, pid, addrs) + +-- | A relay switch with the NAT handlers registered. +newRelaySwitch :: NATConfig -> IO (Switch, PeerId, Multiaddr, RelayState) +newRelaySwitch config = do + (pid, kp) <- mkTestIdentity + sw <- newSwitch pid kp + addTransport sw =<< newTCPTransport + (relayState, _circuitState) <- registerNATHandlers sw config + addrs <- switchListen sw defaultConnectionGater [loopbackAddr] + case addrs of + (a : _) -> pure (sw, pid, a, relayState) + [] -> fail "relay did not bind a listen address" + +-- | Send RESERVE to the relay over an existing connection and return the +-- status the relay replied with. +reserveOn :: Switch -> Connection -> IO (Maybe RelayStatus) +reserveOn sw conn = do + stream <- newStream sw conn >>= either (fail . show) pure + negotiated <- negotiateInitiator stream [hopProtocolId] + case negotiated of + NoProtocol -> fail "relay does not support the hop protocol" + Accepted _ -> hopStatus <$> (makeReservation stream >>= either fail pure) + +-- | Poll the relay's pool until it holds at least @n@ connections for the +-- peer. Inbound connections are admitted on the accept-loop thread, so a +-- dial returning does not yet mean the relay has pooled its side. +waitForConns :: Switch -> PeerId -> Int -> IO [Connection] +waitForConns sw pid n = go (200 :: Int) + where + go 0 = fail $ "relay never pooled " ++ show n ++ " connection(s) for the peer" + go k = do + conns <- atomically $ lookupAllConns (swConnPool sw) pid + if length conns >= n + then pure conns + else threadDelay 10000 >> go (k - 1) + +-- | Open a second connection to the relay, bypassing the pool. +-- 'LibP2P.Switch.Dial.dial' would return the existing pooled connection, +-- so the raw transport dial and the outbound upgrade are driven directly. +-- Only the relay's view matters here, and the relay pools its own side +-- through the normal accept path. +openSecondConnection :: Switch -> Multiaddr -> IO () +openSecondConnection sw addr = do + transport <- newTCPTransport + rawConn <- transportDial transport addr + _conn <- upgradeOutbound (swIdentityKey sw) rawConn + pure () + +reservedPeers :: RelayState -> IO [PeerId] +reservedPeers relayState = Map.keys <$> readTVarIO (rsReservations relayState) + +spec :: Spec +spec = describe "relay reservation lifecycle" $ do + it "drops the reservation when the reserving peer's last connection closes" $ do + (swR, pidR, addrR, relayState) <- newRelaySwitch defaultNATConfig + (swC, pidC, _) <- newListeningSwitch + connCR <- dial swC pidR [addrR] >>= either (fail . show) pure + status <- reserveOn swC connCR + status `shouldBe` Just RelayOK + reservedPeers relayState `shouldReturn` [pidC] + -- Tear down the relay's own side of the connection: the notifier runs + -- synchronously, so no polling is needed for the assertion. + relayConns <- waitForConns swR pidC 1 + mapM_ (closeConnection swR) relayConns + reservedPeers relayState `shouldReturn` [] + switchClose swC + switchClose swR + + it "frees the reservation slot for another peer without waiting for expiry" $ do + let config = NATConfig { ncRelayConfig = defaultRelayConfig { rcMaxReservations = 1 } } + (swR, pidR, addrR, relayState) <- newRelaySwitch config + (swC1, pidC1, _) <- newListeningSwitch + (swC2, _pidC2, _) <- newListeningSwitch + conn1 <- dial swC1 pidR [addrR] >>= either (fail . show) pure + reserveOn swC1 conn1 `shouldReturn` Just RelayOK + -- The relay is now at capacity + conn2 <- dial swC2 pidR [addrR] >>= either (fail . show) pure + reserveOn swC2 conn2 `shouldReturn` Just ReservationRefused + -- Disconnecting the holder must free the slot immediately + relayConns <- waitForConns swR pidC1 1 + mapM_ (closeConnection swR) relayConns + reservedPeers relayState `shouldReturn` [] + reserveOn swC2 conn2 `shouldReturn` Just RelayOK + switchClose swC1 + switchClose swC2 + switchClose swR + + it "keeps the reservation while another connection to the same peer remains" $ do + (swR, pidR, addrR, relayState) <- newRelaySwitch defaultNATConfig + (swC, pidC, _) <- newListeningSwitch + connCR <- dial swC pidR [addrR] >>= either (fail . show) pure + reserveOn swC connCR `shouldReturn` Just RelayOK + -- A second, independent connection from the same peer. Switch.dial + -- would hand back the pooled one, so this goes straight through the + -- transport and the upgrade pipeline; the relay accepts it as a + -- second inbound connection for pidC. + openSecondConnection swC addrR + relayConns <- waitForConns swR pidC 2 + case relayConns of + (first' : rest@(_ : _)) -> do + -- Losing one connection must not invalidate the reservation + closeConnection swR first' + reservedPeers relayState `shouldReturn` [pidC] + -- Losing the last one must + mapM_ (closeConnection swR) rest + reservedPeers relayState `shouldReturn` [] + _ -> expectationFailure "expected two pooled connections for the peer" + switchClose swC + switchClose swR + + it "is a no-op when the disconnecting peer holds no reservation" $ do + (swR, pidR, addrR, relayState) <- newRelaySwitch defaultNATConfig + (swHolder, pidHolder, _) <- newListeningSwitch + (swPlain, pidPlain, _) <- newListeningSwitch + connHolder <- dial swHolder pidR [addrR] >>= either (fail . show) pure + reserveOn swHolder connHolder `shouldReturn` Just RelayOK + _ <- dial swPlain pidR [addrR] >>= either (fail . show) pure + plainConns <- waitForConns swR pidPlain 1 + mapM_ (closeConnection swR) plainConns + -- The unrelated holder's reservation is untouched + reservedPeers relayState `shouldReturn` [pidHolder] + switchClose swHolder + switchClose swPlain + switchClose swR + + it "does not run the cleanup twice for the same connection" $ do + (swR, pidR, addrR, relayState) <- newRelaySwitch defaultNATConfig + (swC, pidC, _) <- newListeningSwitch + connCR <- dial swC pidR [addrR] >>= either (fail . show) pure + reserveOn swC connCR `shouldReturn` Just RelayOK + relayConns <- waitForConns swR pidC 1 + mapM_ (closeConnection swR) relayConns + reservedPeers relayState `shouldReturn` [] + -- Re-arm a reservation, then tear the same connections down again. + -- closeConnection is idempotent, so the notifier must not fire a + -- second time and must not remove the new reservation. + let rearmed = ActiveReservation { arPeerId = pidC, arExpiration = maxBound } + atomically $ modifyTVar' (rsReservations relayState) (Map.insert pidC rearmed) + mapM_ (closeConnection swR) relayConns + reservedPeers relayState `shouldReturn` [pidC] + switchClose swC + switchClose swR