-- | Ping protocol implementation (specs/ping).
--
-- Protocol ID: /ipfs/ping/1.0.0
--
-- Wire format: 32 bytes random → 32 bytes echo. No framing, no protobuf.
-- The responder runs an echo loop: reads 32 bytes, writes them back,
-- until the initiator closes the stream, then closes its own side.
--
-- The initiator keeps at most one outbound ping stream per peer
-- (ping.md: "The dialing peer MUST NOT keep more than one outbound
-- stream for the ping protocol per peer"). A 'PingSession' holds that
-- single stream and reuses it for successive pings (ping.md: the peer
-- "MAY send further payloads on the same stream"); the stream is closed
-- when the session ends or on the first failed ping. Streams are opened
-- through the Switch ('newStream'), so each session holds exactly one
-- stream reservation, released on close.
--
-- The listener accepts at most two concurrent ping streams per remote
-- peer (ping.md: "The listening peer SHOULD accept at most two streams
-- per peer since cross-stream behavior is non-linear and stream writes
-- occur asynchronously"). 'registerPingHandler' installs a
-- 'PingLimiter' that counts live inbound ping streams per peer and
-- resets the third and subsequent streams without serving them.
module LibP2P.Protocol.Ping
  ( -- * Protocol ID
    pingProtocolId
    -- * Types
  , PingError (..)
  , PingResult (..)
  , PingSession
    -- * Responder
  , handlePing
  , PingLimiter
  , newPingLimiter
  , handlePingLimited
    -- * Initiator
  , sendPing
  , openPingSession
  , ping
  , pingWithTimeout
  , closePingSession
  , withPingSession
    -- * Registration
  , registerPingHandler
    -- * Constants
  , pingSize
  , pingTimeoutMicros
  , maxPingStreamsPerPeer
  ) where

import Control.Concurrent.MVar (MVar, newMVar, withMVar)
import Control.Concurrent.STM
  ( TVar
  , atomically
  , modifyTVar'
  , newTVarIO
  , readTVar
  , writeTVar
  )
import Control.Exception (SomeException, catch, finally, try)
import Control.Monad (unless)
import Data.ByteString (ByteString)
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import Data.Time.Clock (NominalDiffTime, diffUTCTime, getCurrentTime)
import Crypto.Random (getRandomBytes)
import LibP2P.Crypto.PeerId (PeerId)
import LibP2P.MultistreamSelect.Negotiation
  ( StreamIO (..)
  , negotiateInitiator
  , NegotiationResult (..)
  , readExactBounded
  )
import LibP2P.Switch.Connection (newStream)
import LibP2P.Switch.Types
  ( Connection (..)
  , Switch (..)
  )
import System.Timeout (timeout)

-- | Ping protocol ID.
pingProtocolId :: Text
pingProtocolId :: Text
pingProtocolId = Text
"/ipfs/ping/1.0.0"

-- | Ping payload size: 32 bytes.
pingSize :: Int
pingSize :: Int
pingSize = Int
32

-- | Time to wait for an echo before giving up, in microseconds.
-- 10 seconds, mirroring go-libp2p's ping timeout.
pingTimeoutMicros :: Int
pingTimeoutMicros :: Int
pingTimeoutMicros = Int
10000000

-- | Maximum concurrent inbound ping streams served per remote peer
-- (ping.md: "The listening peer SHOULD accept at most two streams per
-- peer since cross-stream behavior is non-linear and stream writes
-- occur asynchronously").
maxPingStreamsPerPeer :: Int
maxPingStreamsPerPeer :: Int
maxPingStreamsPerPeer = Int
2

-- | Ping error types.
data PingError
  = PingTimeout          -- ^ No echo within the timeout
  | PingMismatch         -- ^ Response doesn't match sent bytes
  | PingStreamError !String  -- ^ Stream open, negotiation, or I/O error
  deriving (Int -> PingError -> ShowS
[PingError] -> ShowS
PingError -> [Char]
(Int -> PingError -> ShowS)
-> (PingError -> [Char])
-> ([PingError] -> ShowS)
-> Show PingError
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PingError -> ShowS
showsPrec :: Int -> PingError -> ShowS
$cshow :: PingError -> [Char]
show :: PingError -> [Char]
$cshowList :: [PingError] -> ShowS
showList :: [PingError] -> ShowS
Show, PingError -> PingError -> Bool
(PingError -> PingError -> Bool)
-> (PingError -> PingError -> Bool) -> Eq PingError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PingError -> PingError -> Bool
== :: PingError -> PingError -> Bool
$c/= :: PingError -> PingError -> Bool
/= :: PingError -> PingError -> Bool
Eq)

-- | Successful ping result.
data PingResult = PingResult
  { PingResult -> NominalDiffTime
pingRTT :: !NominalDiffTime  -- ^ Round-trip time
  } deriving (Int -> PingResult -> ShowS
[PingResult] -> ShowS
PingResult -> [Char]
(Int -> PingResult -> ShowS)
-> (PingResult -> [Char])
-> ([PingResult] -> ShowS)
-> Show PingResult
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PingResult -> ShowS
showsPrec :: Int -> PingResult -> ShowS
$cshow :: PingResult -> [Char]
show :: PingResult -> [Char]
$cshowList :: [PingResult] -> ShowS
showList :: [PingResult] -> ShowS
Show, PingResult -> PingResult -> Bool
(PingResult -> PingResult -> Bool)
-> (PingResult -> PingResult -> Bool) -> Eq PingResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PingResult -> PingResult -> Bool
== :: PingResult -> PingResult -> Bool
$c/= :: PingResult -> PingResult -> Bool
/= :: PingResult -> PingResult -> Bool
Eq)

-- | Handle an inbound Ping request (responder / echo loop).
--
-- Reads 32 bytes, writes them back. Repeats until the initiator closes
-- its write side (EOF), then closes this side of the stream (ping.md:
-- the listening peer SHOULD exit the loop and close the stream).
handlePing :: StreamIO -> PeerId -> IO ()
handlePing :: StreamIO -> PeerId -> IO ()
handlePing StreamIO
stream PeerId
_remotePeerId = IO ()
echoLoop IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` StreamIO -> IO ()
closeQuietly StreamIO
stream
  where
    echoLoop :: IO ()
echoLoop = do
      result <- StreamIO -> Int -> Int -> IO (Either [Char] ByteString)
readExactBounded StreamIO
stream Int
pingSize Int
pingSize IO (Either [Char] ByteString)
-> (SomeException -> IO (Either [Char] ByteString))
-> IO (Either [Char] ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch`
                (\(SomeException
_ :: SomeException) -> Either [Char] ByteString -> IO (Either [Char] ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ByteString
forall a b. a -> Either a b
Left [Char]
"stream closed"))
      case result of
        Left [Char]
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()  -- Stream closed, exit loop
        Right ByteString
payload -> do
          StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream ByteString
payload
          IO ()
echoLoop

-- | Per-peer accounting of live inbound ping streams, shared by every
-- invocation of the registered ping handler on one Switch.
newtype PingLimiter = PingLimiter (TVar (Map.Map PeerId Int))

-- | Create an empty inbound ping stream limiter.
newPingLimiter :: IO PingLimiter
newPingLimiter :: IO PingLimiter
newPingLimiter = TVar (Map PeerId Int) -> PingLimiter
PingLimiter (TVar (Map PeerId Int) -> PingLimiter)
-> IO (TVar (Map PeerId Int)) -> IO PingLimiter
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map PeerId Int -> IO (TVar (Map PeerId Int))
forall a. a -> IO (TVar a)
newTVarIO Map PeerId Int
forall k a. Map k a
Map.empty

-- | Serve an inbound ping stream, enforcing the per-peer cap.
--
-- If the remote peer already has 'maxPingStreamsPerPeer' live ping
-- streams, the new stream is reset (closed without serving the echo
-- loop). Otherwise the stream occupies a slot for the duration of
-- 'handlePing'; the slot is released when the stream closes or errors.
handlePingLimited :: PingLimiter -> StreamIO -> PeerId -> IO ()
handlePingLimited :: PingLimiter -> StreamIO -> PeerId -> IO ()
handlePingLimited (PingLimiter TVar (Map PeerId Int)
countsVar) StreamIO
stream PeerId
peer = do
  accepted <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
    counts <- TVar (Map PeerId Int) -> STM (Map PeerId Int)
forall a. TVar a -> STM a
readTVar TVar (Map PeerId Int)
countsVar
    let live = Int -> PeerId -> Map PeerId Int -> Int
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Int
0 PeerId
peer Map PeerId Int
counts
    if live >= maxPingStreamsPerPeer
      then pure False
      else do
        writeTVar countsVar (Map.insert peer (live + 1) counts)
        pure True
  if accepted
    then handlePing stream peer `finally` atomically (modifyTVar' countsVar releaseSlot)
    else closeQuietly stream
  where
    releaseSlot :: Map PeerId Int -> Map PeerId Int
releaseSlot = (Int -> Maybe Int) -> PeerId -> Map PeerId Int -> Map PeerId Int
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (\Int
n -> if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) PeerId
peer

-- | The single outbound ping stream to a peer, negotiated and ready.
--
-- Obtain with 'openPingSession' (or scoped via 'withPingSession'), send
-- pings with 'ping', and always release with 'closePingSession'. A
-- session whose ping failed (timeout, mismatch, I/O error) closes its
-- stream immediately and rejects further pings.
--
-- Concurrent 'ping' calls on one session are serialized on 'psLock':
-- exactly one write/echo exchange runs on the stream at a time, so
-- concurrent callers queue instead of interleaving their 32-byte
-- payloads on the wire.
data PingSession = PingSession
  { PingSession -> StreamIO
psStream :: !StreamIO
  , PingSession -> IORef Bool
psClosed :: !(IORef Bool)
  , PingSession -> MVar ()
psLock   :: !(MVar ())  -- ^ Held for the duration of one ping exchange
  }

-- | Open a ping stream on the connection and negotiate the protocol.
--
-- The stream is opened through the Switch so it is counted against the
-- peer's stream limits; the reservation is released when the session is
-- closed. On any failure the stream (if opened) is closed before
-- returning.
openPingSession :: Switch -> Connection -> IO (Either PingError PingSession)
openPingSession :: Switch -> Connection -> IO (Either PingError PingSession)
openPingSession Switch
sw Connection
conn = do
  streamOrErr <- Switch -> Connection -> IO (Either ResourceError StreamIO)
newStream Switch
sw Connection
conn
  case streamOrErr of
    Left ResourceError
err ->
      Either PingError PingSession -> IO (Either PingError PingSession)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingSession
forall a b. a -> Either a b
Left ([Char] -> PingError
PingStreamError ([Char]
"stream reservation failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ ResourceError -> [Char]
forall a. Show a => a -> [Char]
show ResourceError
err)))
    Right StreamIO
stream -> do
      negotiated <- IO NegotiationResult -> IO (Either SomeException NegotiationResult)
forall e a. Exception e => IO a -> IO (Either e a)
try (StreamIO -> [Text] -> IO NegotiationResult
negotiateInitiator StreamIO
stream [Text
pingProtocolId])
      case negotiated of
        Right (Accepted Text
_) -> do
          closedRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False
          lock <- newMVar ()
          pure (Right (PingSession stream closedRef lock))
        Right NegotiationResult
NoProtocol -> do
          StreamIO -> IO ()
closeQuietly StreamIO
stream
          Either PingError PingSession -> IO (Either PingError PingSession)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingSession
forall a b. a -> Either a b
Left ([Char] -> PingError
PingStreamError [Char]
"remote does not support ping"))
        Left (SomeException
e :: SomeException) -> do
          StreamIO -> IO ()
closeQuietly StreamIO
stream
          Either PingError PingSession -> IO (Either PingError PingSession)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingSession
forall a b. a -> Either a b
Left ([Char] -> PingError
PingStreamError ([Char]
"ping negotiation failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)))

-- | Send one ping on the session with the default timeout
-- ('pingTimeoutMicros'). The session's stream is reused across calls.
ping :: PingSession -> IO (Either PingError PingResult)
ping :: PingSession -> IO (Either PingError PingResult)
ping = Int -> PingSession -> IO (Either PingError PingResult)
pingWithTimeout Int
pingTimeoutMicros

-- | Send one ping on the session, waiting at most the given number of
-- microseconds for the echo. On failure the session is closed: a stream
-- whose echo timed out or went wrong cannot be reused, because a late
-- echo would corrupt the next ping.
--
-- The whole exchange runs under the session lock, so concurrent callers
-- are queued one after another on the single stream. The closed check
-- happens under the lock too: a caller queued behind a failed ping sees
-- the session as closed instead of writing into a poisoned stream.
pingWithTimeout :: Int -> PingSession -> IO (Either PingError PingResult)
pingWithTimeout :: Int -> PingSession -> IO (Either PingError PingResult)
pingWithTimeout Int
timeoutUs PingSession
sess = MVar ()
-> (() -> IO (Either PingError PingResult))
-> IO (Either PingError PingResult)
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar (PingSession -> MVar ()
psLock PingSession
sess) ((() -> IO (Either PingError PingResult))
 -> IO (Either PingError PingResult))
-> (() -> IO (Either PingError PingResult))
-> IO (Either PingError PingResult)
forall a b. (a -> b) -> a -> b
$ \() -> do
  closed <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (PingSession -> IORef Bool
psClosed PingSession
sess)
  if closed
    then pure (Left (PingStreamError "ping session is closed"))
    else do
      result <- pingOnce timeoutUs (psStream sess)
      case result of
        Left PingError
err -> do
          PingSession -> IO ()
closePingSession PingSession
sess
          Either PingError PingResult -> IO (Either PingError PingResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingResult
forall a b. a -> Either a b
Left PingError
err)
        Either PingError PingResult
ok -> Either PingError PingResult -> IO (Either PingError PingResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Either PingError PingResult
ok

-- | One ping exchange on an already-negotiated stream.
pingOnce :: Int -> StreamIO -> IO (Either PingError PingResult)
pingOnce :: Int -> StreamIO -> IO (Either PingError PingResult)
pingOnce Int
timeoutUs StreamIO
stream = do
  payload <- Int -> IO ByteString
forall byteArray. ByteArray byteArray => Int -> IO byteArray
forall (m :: * -> *) byteArray.
(MonadRandom m, ByteArray byteArray) =>
Int -> m byteArray
getRandomBytes Int
pingSize :: IO ByteString
  t0 <- getCurrentTime
  -- try must wrap timeout (not the reverse): an inner try would catch
  -- the Timeout exception itself and misreport it as a stream error.
  outcome <- try $ timeout timeoutUs $ do
    streamWrite stream payload
    either fail pure =<< readExactBounded stream pingSize pingSize
  case outcome of
    Left (SomeException
e :: SomeException) ->
      Either PingError PingResult -> IO (Either PingError PingResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingResult
forall a b. a -> Either a b
Left ([Char] -> PingError
PingStreamError ([Char]
"ping I/O failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)))
    Right Maybe ByteString
Nothing -> Either PingError PingResult -> IO (Either PingError PingResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingResult
forall a b. a -> Either a b
Left PingError
PingTimeout)
    Right (Just ByteString
echo)
      | ByteString
echo ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= ByteString
payload -> Either PingError PingResult -> IO (Either PingError PingResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError PingResult
forall a b. a -> Either a b
Left PingError
PingMismatch)
      | Bool
otherwise -> do
          t1 <- IO UTCTime
getCurrentTime
          pure (Right (PingResult (diffUTCTime t1 t0)))

-- | Close the session's stream (signalling EOF to the responder's echo
-- loop) and release its stream reservation. Idempotent.
closePingSession :: PingSession -> IO ()
closePingSession :: PingSession -> IO ()
closePingSession PingSession
sess = do
  alreadyClosed <- IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (PingSession -> IORef Bool
psClosed PingSession
sess) (\Bool
c -> (Bool
True, Bool
c))
  unless alreadyClosed $ closeQuietly (psStream sess)

-- | Run an action with a ping session, closing it afterwards even if
-- the action throws. Returns Left if the session could not be opened.
withPingSession
  :: Switch
  -> Connection
  -> (PingSession -> IO a)
  -> IO (Either PingError a)
withPingSession :: forall a.
Switch
-> Connection -> (PingSession -> IO a) -> IO (Either PingError a)
withPingSession Switch
sw Connection
conn PingSession -> IO a
action = do
  opened <- Switch -> Connection -> IO (Either PingError PingSession)
openPingSession Switch
sw Connection
conn
  case opened of
    Left PingError
err -> Either PingError a -> IO (Either PingError a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PingError -> Either PingError a
forall a b. a -> Either a b
Left PingError
err)
    Right PingSession
sess -> (a -> Either PingError a
forall a b. b -> Either a b
Right (a -> Either PingError a) -> IO a -> IO (Either PingError a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PingSession -> IO a
action PingSession
sess) IO (Either PingError a) -> IO () -> IO (Either PingError a)
forall a b. IO a -> IO b -> IO a
`finally` PingSession -> IO ()
closePingSession PingSession
sess

-- | Send a single Ping to a remote peer (initiator side).
--
-- Convenience wrapper: opens a ping session, pings once, and closes the
-- stream. For repeated pings to the same peer, use 'withPingSession'
-- to reuse one stream instead of opening one per call.
sendPing :: Switch -> Connection -> IO (Either PingError PingResult)
sendPing :: Switch -> Connection -> IO (Either PingError PingResult)
sendPing Switch
sw Connection
conn = (PingError -> Either PingError PingResult)
-> (Either PingError PingResult -> Either PingError PingResult)
-> Either PingError (Either PingError PingResult)
-> Either PingError PingResult
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either PingError -> Either PingError PingResult
forall a b. a -> Either a b
Left Either PingError PingResult -> Either PingError PingResult
forall a. a -> a
id (Either PingError (Either PingError PingResult)
 -> Either PingError PingResult)
-> IO (Either PingError (Either PingError PingResult))
-> IO (Either PingError PingResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Switch
-> Connection
-> (PingSession -> IO (Either PingError PingResult))
-> IO (Either PingError (Either PingError PingResult))
forall a.
Switch
-> Connection -> (PingSession -> IO a) -> IO (Either PingError a)
withPingSession Switch
sw Connection
conn PingSession -> IO (Either PingError PingResult)
ping

-- | Register the Ping handler on the Switch.
--
-- The installed handler shares one 'PingLimiter', so concurrent inbound
-- ping streams are capped at 'maxPingStreamsPerPeer' per remote peer.
registerPingHandler :: Switch -> IO ()
registerPingHandler :: Switch -> IO ()
registerPingHandler Switch
sw = do
  limiter <- IO PingLimiter
newPingLimiter
  atomically $ do
    protos <- readTVar (swProtocols sw)
    let handler Connection
conn StreamIO
stream = PingLimiter -> StreamIO -> PeerId -> IO ()
handlePingLimited PingLimiter
limiter StreamIO
stream (Connection -> PeerId
connPeerId Connection
conn)
    writeTVar (swProtocols sw) (Map.insert pingProtocolId handler protos)

-- | Close a stream, swallowing any exception (best-effort EOF signal).
closeQuietly :: StreamIO -> IO ()
closeQuietly :: StreamIO -> IO ()
closeQuietly StreamIO
stream = StreamIO -> IO ()
streamClose StreamIO
stream IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()