-- | NAT traversal handler registration (specs/autonat, specs/relay, specs/relay/DCUtR).
--
-- Wires the AutoNAT, Circuit Relay v2, and DCUtR module implementations
-- into the Switch protocol registry, in the style of
-- 'LibP2P.Protocol.Identify.registerIdentifyHandlers':
--
--   /libp2p/autonat/1.0.0            — AutoNAT dial-back server
--   /libp2p/circuit/relay/0.2.0/hop  — Circuit Relay v2 relay server
--   /libp2p/circuit/relay/0.2.0/stop — Circuit Relay v2 target (inbound relayed streams)
--   /libp2p/dcutr                    — DCUtR hole-punch coordination (handler side)
module LibP2P.NAT
  ( -- * Configuration
    NATConfig (..)
  , defaultNATConfig
    -- * Registration
  , registerNATHandlers
  , registerAutoNATHandler
  , registerRelayHopHandler
  , registerRelayStopHandler
  , registerDCUtRHandler
  ) where

import Control.Concurrent.STM (atomically)
import Control.Exception (SomeException, catch, try)
import LibP2P.Crypto.PeerId (PeerId, peerIdBytes)
import LibP2P.Multiaddr (Multiaddr (..), encapsulate)
import LibP2P.Multiaddr.Protocol (Protocol (..))
import LibP2P.MultistreamSelect.Negotiation
  ( NegotiationResult (..)
  , StreamIO (..)
  , negotiateInitiator
  )
import LibP2P.NAT.AutoNAT (AutoNATConfig (..), handleAutoNAT)
import LibP2P.NAT.AutoNAT.Message (autoNATProtocolId)
import LibP2P.NAT.DCUtR (DCUtRConfig (..), handleDCUtR)
import LibP2P.NAT.DCUtR.Message (dcutrProtocolId)
import LibP2P.NAT.Relay
  ( HopContext (..)
  , RelayConfig
  , RelayState
  , defaultRelayConfig
  , handleConnect
  , handleReserve
  , newRelayState
  )
import LibP2P.NAT.Relay.Client (handleStop)
import LibP2P.NAT.Relay.Message
  ( HopMessage (..)
  , HopMessageType (..)
  , RelayLimit
  , RelayStatus (..)
  , hopProtocolId
  , maxRelayMessageSize
  , readHopMessage
  , stopProtocolId
  , writeHopMessage
  )
import LibP2P.Switch (selectTransport, setStreamHandler)
import LibP2P.Switch.ConnPool (lookupConn)
import LibP2P.Switch.Connection (newStream)
import LibP2P.Switch.Dial (dial)
import LibP2P.Switch.Listen (switchListenAddrs)
import LibP2P.Switch.Types
  ( Connection (..)
  , MuxerSession (..)
  , Switch (..)
  )
import LibP2P.Switch.Upgrade (upgradeOutbound)
import LibP2P.Transport (Transport (..))

-- | Configuration for the NAT traversal handlers.
data NATConfig = NATConfig
  { NATConfig -> RelayConfig
ncRelayConfig     :: !RelayConfig
    -- ^ Resource limits for the Circuit Relay v2 server side
  , NATConfig -> PeerId -> Maybe RelayLimit -> StreamIO -> IO ()
ncOnRelayedStream :: !(PeerId -> Maybe RelayLimit -> StreamIO -> IO ())
    -- ^ Invoked when a relay delivers an inbound relayed stream (stop
    -- protocol) after the CONNECT/OK exchange: source peer, limit
    -- advertised by the relay, and the relayed stream. The application
    -- owns the stream from this point (e.g. to run DCUtR over it).
  }

-- | Default NAT configuration: default relay limits, and inbound relayed
-- streams are left to the remote end (no local consumer).
defaultNATConfig :: NATConfig
defaultNATConfig :: NATConfig
defaultNATConfig = NATConfig
  { ncRelayConfig :: RelayConfig
ncRelayConfig     = RelayConfig
defaultRelayConfig
  , ncOnRelayedStream :: PeerId -> Maybe RelayLimit -> StreamIO -> IO ()
ncOnRelayedStream = \PeerId
_ Maybe RelayLimit
_ StreamIO
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  }

-- | Register all four NAT protocol handlers on the Switch.
--
-- Creates the relay server state from 'ncRelayConfig' and returns it so
-- callers can inspect reservations/circuits.
registerNATHandlers :: Switch -> NATConfig -> IO RelayState
registerNATHandlers :: Switch -> NATConfig -> IO RelayState
registerNATHandlers Switch
sw NATConfig
config = do
  relayState <- RelayConfig -> IO RelayState
newRelayState (NATConfig -> RelayConfig
ncRelayConfig NATConfig
config)
  registerAutoNATHandler sw
  registerRelayHopHandler sw relayState
  registerRelayStopHandler sw (ncOnRelayedStream config)
  registerDCUtRHandler sw
  pure relayState

-- | Register the AutoNAT server handler (/libp2p/autonat/1.0.0).
--
-- The dial-back deliberately bypasses the connection pool: reusing the
-- requester's existing connection would always report success. Instead a
-- fresh transport dial + upgrade verifies both reachability and identity,
-- and the probe connection is closed immediately (go-libp2p uses a
-- separate dialer host for the same reason).
registerAutoNATHandler :: Switch -> IO ()
registerAutoNATHandler :: Switch -> IO ()
registerAutoNATHandler Switch
sw =
  Switch -> ProtocolId -> StreamHandler -> IO ()
setStreamHandler Switch
sw ProtocolId
autoNATProtocolId (StreamHandler -> IO ()) -> StreamHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn StreamIO
stream ->
    let config :: AutoNATConfig
config = AutoNATConfig
          { natThreshold :: Int
natThreshold = Int
3
          , natDialBack :: PeerId -> [Multiaddr] -> IO (Either [Char] ())
natDialBack  = Switch -> PeerId -> [Multiaddr] -> IO (Either [Char] ())
freshDialBack Switch
sw
          }
    in AutoNATConfig -> StreamIO -> PeerId -> Multiaddr -> IO ()
handleAutoNAT AutoNATConfig
config StreamIO
stream (Connection -> PeerId
connPeerId Connection
conn) (Connection -> Multiaddr
connRemoteAddr Connection
conn)

-- | Dial back a peer on a fresh connection, verify its identity, and close.
freshDialBack :: Switch -> PeerId -> [Multiaddr] -> IO (Either String ())
freshDialBack :: Switch -> PeerId -> [Multiaddr] -> IO (Either [Char] ())
freshDialBack Switch
_ PeerId
_ [] = Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left [Char]
"dial-back: no addresses")
freshDialBack Switch
sw PeerId
pid (Multiaddr
addr : [Multiaddr]
rest) = do
  result <- IO (Either [Char] ())
-> IO (Either SomeException (Either [Char] ()))
forall e a. Exception e => IO a -> IO (Either e a)
try (Switch -> PeerId -> Multiaddr -> IO (Either [Char] ())
probeAddr Switch
sw PeerId
pid Multiaddr
addr)
  case result of
    Right (Right ()) -> Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (() -> Either [Char] ()
forall a b. b -> Either a b
Right ())
    Right (Left [Char]
err)
      | [Multiaddr] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Multiaddr]
rest -> Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left [Char]
err)
      | Bool
otherwise -> Switch -> PeerId -> [Multiaddr] -> IO (Either [Char] ())
freshDialBack Switch
sw PeerId
pid [Multiaddr]
rest
    Left (SomeException
e :: SomeException)
      | [Multiaddr] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Multiaddr]
rest -> Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left (SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e))
      | Bool
otherwise -> Switch -> PeerId -> [Multiaddr] -> IO (Either [Char] ())
freshDialBack Switch
sw PeerId
pid [Multiaddr]
rest

-- | Probe a single address: transport dial, upgrade, check peer identity.
probeAddr :: Switch -> PeerId -> Multiaddr -> IO (Either String ())
probeAddr :: Switch -> PeerId -> Multiaddr -> IO (Either [Char] ())
probeAddr Switch
sw PeerId
pid Multiaddr
addr = do
  mTransport <- Switch -> Multiaddr -> IO (Maybe Transport)
selectTransport Switch
sw Multiaddr
addr
  case mTransport of
    Maybe Transport
Nothing -> Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left ([Char]
"dial-back: no transport for " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Multiaddr -> [Char]
forall a. Show a => a -> [Char]
show Multiaddr
addr))
    Just Transport
transport -> do
      rawConn <- Transport -> Multiaddr -> IO RawConnection
transportDial Transport
transport Multiaddr
addr
      conn <- upgradeOutbound (swIdentityKey sw) rawConn
      let matches = Connection -> PeerId
connPeerId Connection
conn PeerId -> PeerId -> Bool
forall a. Eq a => a -> a -> Bool
== PeerId
pid
      muxClose (connSession conn) `catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      pure $ if matches
        then Right ()
        else Left "dial-back: peer identity mismatch"

-- | Register the Circuit Relay v2 hop handler
-- (/libp2p/circuit/relay/0.2.0/hop): serve RESERVE and CONNECT requests.
registerRelayHopHandler :: Switch -> RelayState -> IO ()
registerRelayHopHandler :: Switch -> RelayState -> IO ()
registerRelayHopHandler Switch
sw RelayState
relayState =
  Switch -> ProtocolId -> StreamHandler -> IO ()
setStreamHandler Switch
sw ProtocolId
hopProtocolId (StreamHandler -> IO ()) -> StreamHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn StreamIO
stream -> do
    result <- StreamIO -> Int -> IO (Either [Char] HopMessage)
readHopMessage StreamIO
stream Int
maxRelayMessageSize
    case result of
      Left [Char]
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Right HopMessage
msg -> case HopMessage -> Maybe HopMessageType
hopType HopMessage
msg of
        Just HopMessageType
HopReserve -> do
          ctx <- Switch -> Connection -> IO HopContext
switchHopContext Switch
sw Connection
conn
          handleReserve relayState ctx stream (connPeerId conn)
        Just HopMessageType
HopConnect -> do
          ctx <- Switch -> Connection -> IO HopContext
switchHopContext Switch
sw Connection
conn
          handleConnect relayState ctx stream (connPeerId conn) msg (openStopStream sw)
        Maybe HopMessageType
_ -> StreamIO -> HopMessage -> IO ()
writeHopMessage StreamIO
stream HopMessage
          { hopType :: Maybe HopMessageType
hopType = HopMessageType -> Maybe HopMessageType
forall a. a -> Maybe a
Just HopMessageType
HopStatus
          , hopPeer :: Maybe RelayPeer
hopPeer = Maybe RelayPeer
forall a. Maybe a
Nothing
          , hopReservation :: Maybe Reservation
hopReservation = Maybe Reservation
forall a. Maybe a
Nothing
          , hopLimit :: Maybe RelayLimit
hopLimit = Maybe RelayLimit
forall a. Maybe a
Nothing
          , hopStatus :: Maybe RelayStatus
hopStatus = RelayStatus -> Maybe RelayStatus
forall a. a -> Maybe a
Just RelayStatus
UnexpectedMessage
          }

-- | Build the per-request hop context from the Switch: the relay's own
-- identity (signs reservation vouchers), its listen addresses with the
-- @/p2p/\<relay\>@ suffix the circuit-v2 spec requires for reservation
-- addrs, and the address the requesting connection arrived over.
switchHopContext :: Switch -> Connection -> IO HopContext
switchHopContext :: Switch -> Connection -> IO HopContext
switchHopContext Switch
sw Connection
conn = do
  addrs <- Switch -> IO [Multiaddr]
switchListenAddrs Switch
sw
  let relayP2P = [Protocol] -> Multiaddr
Multiaddr [ByteString -> Protocol
P2P (PeerId -> ByteString
peerIdBytes (Switch -> PeerId
swLocalPeerId Switch
sw))]
  pure HopContext
    { hcRelayId    = swLocalPeerId sw
    , hcRelayKey   = swIdentityKey sw
    , hcRelayAddrs = map (`encapsulate` relayP2P) addrs
    , hcRemoteAddr = connRemoteAddr conn
    }

-- | Open a stop-protocol stream to the circuit target over an existing
-- connection. Returns Nothing when the target is not connected or the
-- stop protocol cannot be negotiated.
openStopStream :: Switch -> PeerId -> IO (Maybe StreamIO)
openStopStream :: Switch -> PeerId -> IO (Maybe StreamIO)
openStopStream Switch
sw PeerId
targetId = do
  result <- IO (Maybe StreamIO) -> IO (Either SomeException (Maybe StreamIO))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Maybe StreamIO) -> IO (Either SomeException (Maybe StreamIO)))
-> IO (Maybe StreamIO)
-> IO (Either SomeException (Maybe StreamIO))
forall a b. (a -> b) -> a -> b
$ do
    mConn <- STM (Maybe Connection) -> IO (Maybe Connection)
forall a. STM a -> IO a
atomically (STM (Maybe Connection) -> IO (Maybe Connection))
-> STM (Maybe Connection) -> IO (Maybe Connection)
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId [Connection]) -> PeerId -> STM (Maybe Connection)
lookupConn (Switch -> TVar (Map PeerId [Connection])
swConnPool Switch
sw) PeerId
targetId
    case mConn of
      Maybe Connection
Nothing -> Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StreamIO
forall a. Maybe a
Nothing
      Just Connection
conn -> do
        streamOrErr <- Switch -> Connection -> IO (Either ResourceError StreamIO)
newStream Switch
sw Connection
conn
        case streamOrErr of
          Left ResourceError
_ -> Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StreamIO
forall a. Maybe a
Nothing
          Right StreamIO
stream -> do
            negotiated <- StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateInitiator StreamIO
stream [ProtocolId
stopProtocolId]
            case negotiated of
              Accepted ProtocolId
_ -> Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StreamIO -> Maybe StreamIO
forall a. a -> Maybe a
Just StreamIO
stream)
              NegotiationResult
NoProtocol -> do
                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 ()
                Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StreamIO
forall a. Maybe a
Nothing
  case result of
    Left (SomeException
_ :: SomeException) -> Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StreamIO
forall a. Maybe a
Nothing
    Right Maybe StreamIO
mStream -> Maybe StreamIO -> IO (Maybe StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe StreamIO
mStream

-- | Register the Circuit Relay v2 stop handler
-- (/libp2p/circuit/relay/0.2.0/stop): accept inbound relayed streams and
-- hand them to the application callback.
registerRelayStopHandler
  :: Switch
  -> (PeerId -> Maybe RelayLimit -> StreamIO -> IO ())
  -> IO ()
registerRelayStopHandler :: Switch
-> (PeerId -> Maybe RelayLimit -> StreamIO -> IO ()) -> IO ()
registerRelayStopHandler Switch
sw PeerId -> Maybe RelayLimit -> StreamIO -> IO ()
onRelayedStream =
  Switch -> ProtocolId -> StreamHandler -> IO ()
setStreamHandler Switch
sw ProtocolId
stopProtocolId (StreamHandler -> IO ()) -> StreamHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
_conn StreamIO
stream -> do
    result <- StreamIO -> IO (Either [Char] (PeerId, Maybe RelayLimit))
handleStop StreamIO
stream
    case result of
      Left [Char]
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Right (PeerId
sourcePeer, Maybe RelayLimit
mLimit) -> PeerId -> Maybe RelayLimit -> StreamIO -> IO ()
onRelayedStream PeerId
sourcePeer Maybe RelayLimit
mLimit StreamIO
stream

-- | Register the DCUtR handler (/libp2p/dcutr).
--
-- Answers the CONNECT/SYNC exchange with our listen addresses and dials
-- the initiator's addresses through the Switch for the hole punch.
registerDCUtRHandler :: Switch -> IO ()
registerDCUtRHandler :: Switch -> IO ()
registerDCUtRHandler Switch
sw =
  Switch -> ProtocolId -> StreamHandler -> IO ()
setStreamHandler Switch
sw ProtocolId
dcutrProtocolId (StreamHandler -> IO ()) -> StreamHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \Connection
conn StreamIO
stream -> do
    addrs <- Switch -> IO [Multiaddr]
switchListenAddrs Switch
sw
    let config = DCUtRConfig
          { dcMaxAttempts :: Int
dcMaxAttempts = Int
3
          , dcDialer :: Multiaddr -> IO (Either [Char] ())
dcDialer = \Multiaddr
addr -> do
              dialed <- Switch -> PeerId -> [Multiaddr] -> IO (Either DialError Connection)
dial Switch
sw (Connection -> PeerId
connPeerId Connection
conn) [Multiaddr
addr]
              pure $ either (Left . show) (const (Right ())) dialed
          }
    _ <- handleDCUtR config stream addrs
    pure ()