fix: invalidate relay reservations when the peer's last connection closes

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

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

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
     }
 

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).

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
   }

test/LibP2P/DHT/APISpec.hs

@@ -47,7 +47,7 @@
         -- as the closest peer for any lookup
         case msgType msg of
           FindNode -> do
             now <- getCurrentTime
             let selfEntry = DHTPeer
                   { dhtPeerId = peerIdBytes pid
                   , dhtPeerAddrs = []
@@ -126,7 +126,7 @@
         Nothing -> expectationFailure "Expected provider entry"
         Just ents -> do
           length ents `shouldBe` 1
           peProvider (head ents) `shouldBe` localPid
 
     it "sends ADD_PROVIDER messages to closest peers when lookup returns peers" $ do
       sentLog <- newTVarIO []
@@ -198,6 +198,7 @@
   resMgr <- mkMockResourceMgr
   peerStore <- newTVarIO Map.empty
   notifiers <- newTVarIO []
+  disconnectNotifiers <- newTVarIO []
   listeners <- newTVarIO []
   kp <- getDummyKeyPair
   pure Switch
@@ -213,6 +214,7 @@
     , swResourceMgr  = resMgr
     , swPeerStore    = peerStore
     , swNotifiers    = notifiers
+    , swDisconnectNotifiers = disconnectNotifiers
     , swListeners    = listeners
     }
 

test/LibP2P/DHT/DHTSpec.hs

@@ -70,6 +70,7 @@
   resMgr <- mkMockResourceMgr
   peerStore <- newTVarIO Map.empty
   notifiers <- newTVarIO []
+  disconnectNotifiers <- newTVarIO []
   listeners <- newTVarIO []
   pure Switch
     { swLocalPeerId  = pid
@@ -84,6 +85,7 @@
     , swResourceMgr  = resMgr
     , swPeerStore    = peerStore
     , swNotifiers    = notifiers
+    , swDisconnectNotifiers = disconnectNotifiers
     , swListeners    = listeners
     }
 
@@ -425,7 +427,7 @@
       -- Provider should now be persisted
       stored <- getProviders node key
       length stored `shouldBe` 1
       peProvider (head stored) `shouldBe` remotePid
 
     it "ADD_PROVIDER round-trip via GET_PROVIDERS" $ do
       node <- mkTestNode localPid
@@ -459,7 +461,7 @@
         Right resp -> do
           msgType resp `shouldBe` GetProviders
           length (msgProviderPeers resp) `shouldBe` 1
           dhtPeerId (head (msgProviderPeers resp)) `shouldBe` peerIdBytes remotePid
         Left err -> expectationFailure $ "Failed: " ++ err
 
     -- Issue #147: specs/kad-dht requires handling additional RPC request
@@ -753,7 +755,7 @@
           bucketPeers = sameBucketPeers (kValue + 1)
           initial = take kValue bucketPeers
           newcomer = last bucketPeers
           lrs = head initial
       mapM_ (\pid -> do
                e <- mkEntry pid
                atomically $ modifyTVar' (dhtRoutingTable node) (fst . insertPeer e))
@@ -774,7 +776,7 @@
           bucketPeers = sameBucketPeers (kValue + 1)
           initial = take kValue bucketPeers
           newcomer = last bucketPeers
           lrs = head initial
       mapM_ (\pid -> do
                e <- mkEntry pid
                atomically $ modifyTVar' (dhtRoutingTable node) (fst . insertPeer e))

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