-- | DCUtR (Direct Connection Upgrade through Relay) protocol.
--
-- Protocol: /libp2p/dcutr
-- Coordinates hole punching over a relayed connection using a 3-message exchange
-- with RTT-based timing synchronization.
--
-- Message flow:
--   B (initiator) sends CONNECT with B's observed addresses
--   A (handler) sends CONNECT with A's observed addresses
--   B sends SYNC
--   B waits RTT/2, then dials A's addresses
--   A receives SYNC, then dials B's addresses immediately
--   Both peers attempt direct connections at approximately the same time
--
-- Per the spec, every exchanged address is dialled in parallel (hole punching
-- depends on near-simultaneous packets), and on failure the whole exchange is
-- re-run from the CONNECT step so RTT is re-measured (3 attempts total).
module LibP2P.NAT.DCUtR
  ( -- * Types
    DCUtRConfig (..)
  , DCUtRResult (..)
    -- * Protocol operations
  , initiateDCUtR
  , handleDCUtR
    -- * Variants for testing
  , initiateDCUtRWithRTT
  , initiateDCUtRCapture
  , handleDCUtRCapture
  ) where

import qualified Data.ByteString as BS
import Data.IORef (IORef, newIORef, writeIORef)
import Control.Concurrent (threadDelay)
import Control.Concurrent.Async (async, waitAny, cancel)
import Control.Exception (bracket)
import Data.Time.Clock (getCurrentTime, diffUTCTime, NominalDiffTime)
import LibP2P.NAT.DCUtR.Message
import LibP2P.MultistreamSelect.Negotiation (StreamIO (..))
import LibP2P.Multiaddr (Multiaddr, toBytes, fromBytes)

-- | DCUtR configuration.
data DCUtRConfig = DCUtRConfig
  { DCUtRConfig -> Int
dcMaxAttempts :: !Int
    -- ^ Total number of hole punch attempts; the full CONNECT/SYNC exchange
    -- is re-run for each attempt so RTT is re-measured (spec: 3 total =
    -- 1 initial + 2 retries)
  , DCUtRConfig -> Multiaddr -> IO (Either [Char] ())
dcDialer     :: !(Multiaddr -> IO (Either String ()))
    -- ^ Injectable dial function for testing
  }

-- | DCUtR result.
data DCUtRResult = DCUtRSuccess | DCUtRFailed String
  deriving (Int -> DCUtRResult -> ShowS
[DCUtRResult] -> ShowS
DCUtRResult -> [Char]
(Int -> DCUtRResult -> ShowS)
-> (DCUtRResult -> [Char])
-> ([DCUtRResult] -> ShowS)
-> Show DCUtRResult
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DCUtRResult -> ShowS
showsPrec :: Int -> DCUtRResult -> ShowS
$cshow :: DCUtRResult -> [Char]
show :: DCUtRResult -> [Char]
$cshowList :: [DCUtRResult] -> ShowS
showList :: [DCUtRResult] -> ShowS
Show, DCUtRResult -> DCUtRResult -> Bool
(DCUtRResult -> DCUtRResult -> Bool)
-> (DCUtRResult -> DCUtRResult -> Bool) -> Eq DCUtRResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DCUtRResult -> DCUtRResult -> Bool
== :: DCUtRResult -> DCUtRResult -> Bool
$c/= :: DCUtRResult -> DCUtRResult -> Bool
/= :: DCUtRResult -> DCUtRResult -> Bool
Eq)

-- | Outcome of a single hole punch attempt. Only dial failures are
-- retryable; protocol errors abort the upgrade immediately.
data AttemptError = FatalError String | DialError String

-- | Peer B (initiator): run the DCUtR exchange over a relayed stream.
--
-- Flow (repeated up to 'dcMaxAttempts' times while the hole punch fails):
--   1. Send CONNECT with own observed addresses
--   2. Read A's CONNECT (measure RTT)
--   3. Send SYNC
--   4. Wait RTT/2, then dial all of A's addresses in parallel
initiateDCUtR :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IO DCUtRResult
initiateDCUtR :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IO DCUtRResult
initiateDCUtR DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs = do
  rttRef <- Maybe NominalDiffTime -> IO (IORef (Maybe NominalDiffTime))
forall a. a -> IO (IORef a)
newIORef Maybe NominalDiffTime
forall a. Maybe a
Nothing
  initiateDCUtRWithRTT config stream addrs rttRef

-- | Initiator variant that captures RTT for testing.
initiateDCUtRWithRTT :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IORef (Maybe NominalDiffTime) -> IO DCUtRResult
initiateDCUtRWithRTT :: DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> IORef (Maybe NominalDiffTime)
-> IO DCUtRResult
initiateDCUtRWithRTT DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs IORef (Maybe NominalDiffTime)
rttRef =
  DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts DCUtRConfig
config (DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef (Maybe NominalDiffTime))
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
initiatorAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs (IORef (Maybe NominalDiffTime)
-> Maybe (IORef (Maybe NominalDiffTime))
forall a. a -> Maybe a
Just IORef (Maybe NominalDiffTime)
rttRef) Maybe (IORef [ByteString])
forall a. Maybe a
Nothing)

-- | Initiator variant that captures received addresses for testing.
initiateDCUtRCapture :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IORef [BS.ByteString] -> IO DCUtRResult
initiateDCUtRCapture :: DCUtRConfig
-> StreamIO -> [Multiaddr] -> IORef [ByteString] -> IO DCUtRResult
initiateDCUtRCapture DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs IORef [ByteString]
receivedRef =
  DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts DCUtRConfig
config (DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef (Maybe NominalDiffTime))
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
initiatorAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs Maybe (IORef (Maybe NominalDiffTime))
forall a. Maybe a
Nothing (IORef [ByteString] -> Maybe (IORef [ByteString])
forall a. a -> Maybe a
Just IORef [ByteString]
receivedRef))

-- | Peer A (handler): handle the DCUtR exchange over a relayed stream.
--
-- Flow (repeated up to 'dcMaxAttempts' times while the hole punch fails,
-- matching the initiator's retries of the exchange):
--   1. Read B's CONNECT
--   2. Send CONNECT with own observed addresses
--   3. Read SYNC
--   4. Dial all of B's addresses in parallel immediately
handleDCUtR :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IO DCUtRResult
handleDCUtR :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IO DCUtRResult
handleDCUtR DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs =
  DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts DCUtRConfig
config (DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
handlerAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs Maybe (IORef [ByteString])
forall a. Maybe a
Nothing)

-- | Handler variant that captures received addresses for testing.
handleDCUtRCapture :: DCUtRConfig -> StreamIO -> [Multiaddr] -> IORef [BS.ByteString] -> IO DCUtRResult
handleDCUtRCapture :: DCUtRConfig
-> StreamIO -> [Multiaddr] -> IORef [ByteString] -> IO DCUtRResult
handleDCUtRCapture DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs IORef [ByteString]
receivedRef =
  DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts DCUtRConfig
config (DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
handlerAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs (IORef [ByteString] -> Maybe (IORef [ByteString])
forall a. a -> Maybe a
Just IORef [ByteString]
receivedRef))

-- Attempt loop

-- | Run hole punch attempts until one succeeds, a protocol error occurs, or
-- the attempt budget is exhausted.
runAttempts :: DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts :: DCUtRConfig -> IO (Either AttemptError ()) -> IO DCUtRResult
runAttempts DCUtRConfig
config IO (Either AttemptError ())
attempt = Int -> IO DCUtRResult
go Int
1
  where
    maxAttempts :: Int
maxAttempts = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (DCUtRConfig -> Int
dcMaxAttempts DCUtRConfig
config)
    go :: Int -> IO DCUtRResult
go Int
n = do
      result <- IO (Either AttemptError ())
attempt
      case result of
        Right () -> DCUtRResult -> IO DCUtRResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DCUtRResult
DCUtRSuccess
        Left (FatalError [Char]
err) -> DCUtRResult -> IO DCUtRResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> DCUtRResult
DCUtRFailed [Char]
err)
        Left (DialError [Char]
err)
          | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxAttempts -> DCUtRResult -> IO DCUtRResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> DCUtRResult
DCUtRFailed [Char]
err)
          | Bool
otherwise -> Int -> IO DCUtRResult
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | One initiator attempt: full CONNECT/CONNECT/SYNC exchange followed by the
-- synchronized parallel dial.
initiatorAttempt
  :: DCUtRConfig
  -> StreamIO
  -> [Multiaddr]
  -> Maybe (IORef (Maybe NominalDiffTime))
  -> Maybe (IORef [BS.ByteString])
  -> IO (Either AttemptError ())
initiatorAttempt :: DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef (Maybe NominalDiffTime))
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
initiatorAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs Maybe (IORef (Maybe NominalDiffTime))
mRttRef Maybe (IORef [ByteString])
mCaptureRef = do
  let connectOut :: HolePunchMessage
connectOut = HolePunchMessage { hpType :: HolePunchType
hpType = HolePunchType
HPConnect, hpObsAddrs :: [ByteString]
hpObsAddrs = (Multiaddr -> ByteString) -> [Multiaddr] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Multiaddr -> ByteString
toBytes [Multiaddr]
addrs }
  -- Step 1: Send CONNECT with our observed addresses. The spec starts the
  -- RTT timer when the CONNECT is sent, so take t0 before the (possibly
  -- blocking) relayed write.
  t0 <- IO UTCTime
getCurrentTime
  writeHolePunchMessage stream connectOut
  -- Step 2: Read A's CONNECT response (this measures RTT)
  result <- readHolePunchMessage stream maxDCUtRMessageSize
  case result of
    Left [Char]
err -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError ([Char]
"failed to read CONNECT: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
err)))
    Right HolePunchMessage
msg
      | HolePunchMessage -> HolePunchType
hpType HolePunchMessage
msg HolePunchType -> HolePunchType -> Bool
forall a. Eq a => a -> a -> Bool
/= HolePunchType
HPConnect -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError [Char]
"expected CONNECT message"))
      | Bool
otherwise -> do
          t1 <- IO UTCTime
getCurrentTime
          let rtt = UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
t1 UTCTime
t0
          mapM_ (`writeIORef` Just rtt) mRttRef
          mapM_ (`writeIORef` hpObsAddrs msg) mCaptureRef
          -- Step 3: Send SYNC
          writeHolePunchMessage stream (HolePunchMessage { hpType = HPSync, hpObsAddrs = [] })
          -- Step 4: Wait RTT/2, then dial all of A's addresses in parallel
          threadDelay (max 0 (round (rtt * 1000000 / 2)))
          dialAllConcurrently config (parseAddrs (hpObsAddrs msg))

-- | One handler attempt: answer the CONNECT/SYNC exchange, then dial all of
-- the initiator's addresses in parallel.
handlerAttempt
  :: DCUtRConfig
  -> StreamIO
  -> [Multiaddr]
  -> Maybe (IORef [BS.ByteString])
  -> IO (Either AttemptError ())
handlerAttempt :: DCUtRConfig
-> StreamIO
-> [Multiaddr]
-> Maybe (IORef [ByteString])
-> IO (Either AttemptError ())
handlerAttempt DCUtRConfig
config StreamIO
stream [Multiaddr]
addrs Maybe (IORef [ByteString])
mCaptureRef = do
  -- Step 1: Read B's CONNECT
  result <- StreamIO -> Int -> IO (Either [Char] HolePunchMessage)
readHolePunchMessage StreamIO
stream Int
maxDCUtRMessageSize
  case result of
    Left [Char]
err -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError ([Char]
"failed to read CONNECT: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
err)))
    Right HolePunchMessage
msg
      | HolePunchMessage -> HolePunchType
hpType HolePunchMessage
msg HolePunchType -> HolePunchType -> Bool
forall a. Eq a => a -> a -> Bool
/= HolePunchType
HPConnect -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError [Char]
"expected CONNECT message"))
      | Bool
otherwise -> do
          (IORef [ByteString] -> IO ())
-> Maybe (IORef [ByteString]) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (IORef [ByteString] -> [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
`writeIORef` HolePunchMessage -> [ByteString]
hpObsAddrs HolePunchMessage
msg) Maybe (IORef [ByteString])
mCaptureRef
          -- Step 2: Send our CONNECT response
          StreamIO -> HolePunchMessage -> IO ()
writeHolePunchMessage StreamIO
stream (HolePunchMessage { hpType :: HolePunchType
hpType = HolePunchType
HPConnect, hpObsAddrs :: [ByteString]
hpObsAddrs = (Multiaddr -> ByteString) -> [Multiaddr] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Multiaddr -> ByteString
toBytes [Multiaddr]
addrs })
          -- Step 3: Read SYNC
          syncResult <- StreamIO -> Int -> IO (Either [Char] HolePunchMessage)
readHolePunchMessage StreamIO
stream Int
maxDCUtRMessageSize
          case syncResult of
            Left [Char]
err -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError ([Char]
"failed to read SYNC: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
err)))
            Right HolePunchMessage
syncMsg
              | HolePunchMessage -> HolePunchType
hpType HolePunchMessage
syncMsg HolePunchType -> HolePunchType -> Bool
forall a. Eq a => a -> a -> Bool
/= HolePunchType
HPSync -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
FatalError [Char]
"expected SYNC message"))
              | Bool
otherwise ->
                  -- Step 4: Dial all of B's addresses in parallel immediately
                  DCUtRConfig -> [Multiaddr] -> IO (Either AttemptError ())
dialAllConcurrently DCUtRConfig
config ([ByteString] -> [Multiaddr]
parseAddrs (HolePunchMessage -> [ByteString]
hpObsAddrs HolePunchMessage
msg))

-- Helpers

-- | Parse binary multiaddr bytes into Multiaddrs, skipping invalid ones.
parseAddrs :: [BS.ByteString] -> [Multiaddr]
parseAddrs :: [ByteString] -> [Multiaddr]
parseAddrs = (ByteString -> [Multiaddr] -> [Multiaddr])
-> [Multiaddr] -> [ByteString] -> [Multiaddr]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ByteString
bs [Multiaddr]
acc -> case ByteString -> Either [Char] Multiaddr
fromBytes ByteString
bs of Right Multiaddr
a -> Multiaddr
a Multiaddr -> [Multiaddr] -> [Multiaddr]
forall a. a -> [a] -> [a]
: [Multiaddr]
acc; Left [Char]
_ -> [Multiaddr]
acc) []

-- | Dial all addresses concurrently. The first successful dial wins and the
-- remaining dials are cancelled; if every dial fails the attempt is a
-- retryable 'DialError'.
dialAllConcurrently :: DCUtRConfig -> [Multiaddr] -> IO (Either AttemptError ())
dialAllConcurrently :: DCUtRConfig -> [Multiaddr] -> IO (Either AttemptError ())
dialAllConcurrently DCUtRConfig
_config [] = Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
DialError [Char]
"no addresses to dial"))
dialAllConcurrently DCUtRConfig
config [Multiaddr]
addrs =
  IO [Async (Either [Char] ())]
-> ([Async (Either [Char] ())] -> IO ())
-> ([Async (Either [Char] ())] -> IO (Either AttemptError ()))
-> IO (Either AttemptError ())
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket ((Multiaddr -> IO (Async (Either [Char] ())))
-> [Multiaddr] -> IO [Async (Either [Char] ())]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (IO (Either [Char] ()) -> IO (Async (Either [Char] ()))
forall a. IO a -> IO (Async a)
async (IO (Either [Char] ()) -> IO (Async (Either [Char] ())))
-> (Multiaddr -> IO (Either [Char] ()))
-> Multiaddr
-> IO (Async (Either [Char] ()))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DCUtRConfig -> Multiaddr -> IO (Either [Char] ())
dcDialer DCUtRConfig
config) [Multiaddr]
addrs) ((Async (Either [Char] ()) -> IO ())
-> [Async (Either [Char] ())] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Async (Either [Char] ()) -> IO ()
forall a. Async a -> IO ()
cancel) [Async (Either [Char] ())] -> IO (Either AttemptError ())
forall {a}. [Async (Either a ())] -> IO (Either AttemptError ())
waitFirstSuccess
  where
    waitFirstSuccess :: [Async (Either a ())] -> IO (Either AttemptError ())
waitFirstSuccess [] = Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (AttemptError -> Either AttemptError ()
forall a b. a -> Either a b
Left ([Char] -> AttemptError
DialError [Char]
"all dial attempts failed"))
    waitFirstSuccess [Async (Either a ())]
pending = do
      (finished, result) <- [Async (Either a ())] -> IO (Async (Either a ()), Either a ())
forall a. [Async a] -> IO (Async a, a)
waitAny [Async (Either a ())]
pending
      case result of
        Right () -> Either AttemptError () -> IO (Either AttemptError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either AttemptError ()
forall a b. b -> Either a b
Right ())
        Left a
_err -> [Async (Either a ())] -> IO (Either AttemptError ())
waitFirstSuccess ((Async (Either a ()) -> Bool)
-> [Async (Either a ())] -> [Async (Either a ())]
forall a. (a -> Bool) -> [a] -> [a]
filter (Async (Either a ()) -> Async (Either a ()) -> Bool
forall a. Eq a => a -> a -> Bool
/= Async (Either a ())
finished) [Async (Either a ())]
pending)