-- | TCP transport implementation.
--
-- Provides dial/listen for /ip4/.../tcp/... and /ip6/.../tcp/... multiaddrs.
-- Uses the network library for socket operations.
module LibP2P.Transport.TCP
  ( newTCPTransport
  , multiaddrToHostPort
  , socketToStreamIO
  ) where

import Data.Bits (shiftL, shiftR, (.&.), (.|.))
import qualified Data.ByteString as BS
import Data.IP (IPv6, fromHostAddress6, toHostAddress6)
import Data.Word (Word16, Word32, Word8)
import LibP2P.Multiaddr (Multiaddr (..))
import LibP2P.Multiaddr.Protocol (Protocol (..))
import LibP2P.MultistreamSelect.Negotiation (StreamIO (..))
import LibP2P.Transport (Listener (..), RawConnection (..), Transport (..))
import qualified Network.Socket as NS
import qualified Network.Socket.ByteString as NSB

-- | Create a new TCP transport.
newTCPTransport :: IO Transport
newTCPTransport :: IO Transport
newTCPTransport = Transport -> IO Transport
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Transport -> IO Transport) -> Transport -> IO Transport
forall a b. (a -> b) -> a -> b
$ Transport
  { transportDial :: Multiaddr -> IO RawConnection
transportDial = Multiaddr -> IO RawConnection
tcpDial
  , transportListen :: Multiaddr -> IO Listener
transportListen = Multiaddr -> IO Listener
tcpListen
  , transportCanDial :: Multiaddr -> Bool
transportCanDial = Multiaddr -> Bool
canDialTCP
  }

-- | Drop a trailing /p2p/<peer-id> component if present. Identify- and
-- DHT-learned addresses carry the peer ID suffix; the TCP transport
-- connects to the transport portion only.
stripP2P :: Multiaddr -> Multiaddr
stripP2P :: Multiaddr -> Multiaddr
stripP2P (Multiaddr [Protocol]
ps) = case [Protocol] -> [Protocol]
forall a. [a] -> [a]
reverse [Protocol]
ps of
  (P2P ByteString
_ : [Protocol]
rest) -> [Protocol] -> Multiaddr
Multiaddr ([Protocol] -> [Protocol]
forall a. [a] -> [a]
reverse [Protocol]
rest)
  [Protocol]
_ -> [Protocol] -> Multiaddr
Multiaddr [Protocol]
ps

-- | Check if a multiaddr is a TCP address (/ip4/.../tcp/... or
-- /ip6/.../tcp/...), optionally with a trailing /p2p/<peer-id> component.
canDialTCP :: Multiaddr -> Bool
canDialTCP :: Multiaddr -> Bool
canDialTCP Multiaddr
addr = case Multiaddr -> Multiaddr
stripP2P Multiaddr
addr of
  Multiaddr [IP4 Word32
_, TCP Word16
_] -> Bool
True
  Multiaddr [IP6 ByteString
_, TCP Word16
_] -> Bool
True
  Multiaddr
_ -> Bool
False

-- | Extract (HostName, ServiceName) from a TCP multiaddr.
-- A trailing /p2p/<peer-id> component is ignored.
multiaddrToHostPort :: Multiaddr -> Either String (String, String)
multiaddrToHostPort :: Multiaddr -> Either String (String, String)
multiaddrToHostPort Multiaddr
addr = case Multiaddr -> Multiaddr
stripP2P Multiaddr
addr of
  Multiaddr [IP4 Word32
w, TCP Word16
port] -> (String, String) -> Either String (String, String)
forall a b. b -> Either a b
Right (Word32 -> String
renderIPv4 Word32
w, Word16 -> String
forall a. Show a => a -> String
show Word16
port)
  Multiaddr [IP6 ByteString
bs, TCP Word16
port] -> (String, String) -> Either String (String, String)
forall a b. b -> Either a b
Right (ByteString -> String
renderIPv6 ByteString
bs, Word16 -> String
forall a. Show a => a -> String
show Word16
port)
  Multiaddr
_ -> String -> Either String (String, String)
forall a b. a -> Either a b
Left String
"multiaddrToHostPort: expected /ip4/.../tcp/... or /ip6/.../tcp/..."

-- | Dial a TCP address by directly constructing a SockAddr from the Multiaddr.
-- A trailing /p2p/<peer-id> component is stripped before connecting; the
-- original (unstripped) multiaddr is kept as the connection's remote address.
tcpDial :: Multiaddr -> IO RawConnection
tcpDial :: Multiaddr -> IO RawConnection
tcpDial Multiaddr
addr = case Multiaddr -> Multiaddr
stripP2P Multiaddr
addr of
  Multiaddr [IP4 Word32
w, TCP Word16
port] -> do
    let hostAddr :: Word32
hostAddr = (Word8, Word8, Word8, Word8) -> Word32
NS.tupleToHostAddress (Int -> Word32 -> Word8
octet Int
3 Word32
w, Int -> Word32 -> Word8
octet Int
2 Word32
w, Int -> Word32 -> Word8
octet Int
1 Word32
w, Int -> Word32 -> Word8
octet Int
0 Word32
w)
        sockAddr :: SockAddr
sockAddr = PortNumber -> Word32 -> SockAddr
NS.SockAddrInet (Word16 -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
port) Word32
hostAddr
    sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket Family
NS.AF_INET SocketType
NS.Stream ProtocolNumber
NS.defaultProtocol
    NS.connect sock sockAddr
    mkRawConnection sock addr
  Multiaddr [IP6 ByteString
bs, TCP Word16
port] -> do
    let ipv6 :: IPv6
ipv6 = ByteString -> IPv6
bytesToIPv6 ByteString
bs
        hostAddr6 :: HostAddress6
hostAddr6 = IPv6 -> HostAddress6
toHostAddress6 IPv6
ipv6
        sockAddr :: SockAddr
sockAddr = PortNumber -> Word32 -> HostAddress6 -> Word32 -> SockAddr
NS.SockAddrInet6 (Word16 -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
port) Word32
0 HostAddress6
hostAddr6 Word32
0
    sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket Family
NS.AF_INET6 SocketType
NS.Stream ProtocolNumber
NS.defaultProtocol
    NS.connect sock sockAddr
    mkRawConnection sock addr
  Multiaddr
_ -> String -> IO RawConnection
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"tcpDial: unsupported multiaddr"

-- | Create a RawConnection from a connected socket.
--
-- TCP_NODELAY is set on every connection (dialed and accepted), matching
-- go-libp2p and rust-libp2p. libp2p traffic is already coalesced into
-- frames by the muxer, and the many small write-write-read exchanges of
-- the upgrade pipeline (multistream-select, Noise, per-stream
-- negotiation) otherwise stall on Nagle/delayed-ACK interaction —
-- observed as ~270ms per ping open/close cycle in the soak tier.
mkRawConnection :: NS.Socket -> Multiaddr -> IO RawConnection
mkRawConnection :: Socket -> Multiaddr -> IO RawConnection
mkRawConnection Socket
sock Multiaddr
remoteAddr = do
  Socket -> SocketOption -> Int -> IO ()
NS.setSocketOption Socket
sock SocketOption
NS.NoDelay Int
1
  localSockAddr <- Socket -> IO SockAddr
NS.getSocketName Socket
sock
  localAddr <- sockAddrToMultiaddr localSockAddr
  pure RawConnection
    { rcStreamIO = socketToStreamIO sock
    , rcLocalAddr = localAddr
    , rcRemoteAddr = remoteAddr
    , rcClose = NS.close sock
    }

-- | Listen on a TCP address.
tcpListen :: Multiaddr -> IO Listener
tcpListen :: Multiaddr -> IO Listener
tcpListen (Multiaddr [IP4 Word32
w, TCP Word16
port]) = do
  let hostAddr :: Word32
hostAddr = (Word8, Word8, Word8, Word8) -> Word32
NS.tupleToHostAddress (Int -> Word32 -> Word8
octet Int
3 Word32
w, Int -> Word32 -> Word8
octet Int
2 Word32
w, Int -> Word32 -> Word8
octet Int
1 Word32
w, Int -> Word32 -> Word8
octet Int
0 Word32
w)
      sockAddr :: SockAddr
sockAddr = PortNumber -> Word32 -> SockAddr
NS.SockAddrInet (Word16 -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
port) Word32
hostAddr
  sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket Family
NS.AF_INET SocketType
NS.Stream ProtocolNumber
NS.defaultProtocol
  NS.setSocketOption sock NS.ReuseAddr 1
  NS.bind sock sockAddr
  NS.listen sock 256
  boundSockAddr <- NS.getSocketName sock
  boundAddr <- sockAddrToMultiaddr boundSockAddr
  pure Listener
    { listenerAccept = do
        (clientSock, clientSockAddr) <- NS.accept sock
        clientAddr <- sockAddrToMultiaddr clientSockAddr
        mkRawConnection clientSock clientAddr
    , listenerClose = NS.close sock
    , listenerAddr = boundAddr
    }
tcpListen (Multiaddr [IP6 ByteString
bs, TCP Word16
port]) = do
  let ipv6 :: IPv6
ipv6 = ByteString -> IPv6
bytesToIPv6 ByteString
bs
      hostAddr6 :: HostAddress6
hostAddr6 = IPv6 -> HostAddress6
toHostAddress6 IPv6
ipv6
      sockAddr :: SockAddr
sockAddr = PortNumber -> Word32 -> HostAddress6 -> Word32 -> SockAddr
NS.SockAddrInet6 (Word16 -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
port) Word32
0 HostAddress6
hostAddr6 Word32
0
  sock <- Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket Family
NS.AF_INET6 SocketType
NS.Stream ProtocolNumber
NS.defaultProtocol
  NS.setSocketOption sock NS.ReuseAddr 1
  NS.bind sock sockAddr
  NS.listen sock 256
  boundSockAddr <- NS.getSocketName sock
  boundAddr <- sockAddrToMultiaddr boundSockAddr
  pure Listener
    { listenerAccept = do
        (clientSock, clientSockAddr) <- NS.accept sock
        clientAddr <- sockAddrToMultiaddr clientSockAddr
        mkRawConnection clientSock clientAddr
    , listenerClose = NS.close sock
    , listenerAddr = boundAddr
    }
tcpListen Multiaddr
_ = String -> IO Listener
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"tcpListen: unsupported multiaddr"

-- | Convert a Socket to StreamIO.
socketToStreamIO :: NS.Socket -> StreamIO
socketToStreamIO :: Socket -> StreamIO
socketToStreamIO Socket
sock = StreamIO
  { streamWrite :: ByteString -> IO ()
streamWrite = Socket -> ByteString -> IO ()
NSB.sendAll Socket
sock
  , streamReadByte :: IO Word8
streamReadByte = do
      bs <- Socket -> Int -> IO ByteString
NSB.recv Socket
sock Int
1
      if BS.null bs
        then fail "socketToStreamIO: connection closed"
        else pure (BS.head bs)
  , streamClose :: IO ()
streamClose = Socket -> IO ()
NS.close Socket
sock
  }

-- | Convert a SockAddr to a Multiaddr.
sockAddrToMultiaddr :: NS.SockAddr -> IO Multiaddr
sockAddrToMultiaddr :: SockAddr -> IO Multiaddr
sockAddrToMultiaddr (NS.SockAddrInet PortNumber
port Word32
host) = do
  let (Word8
a, Word8
b, Word8
c, Word8
d) = Word32 -> (Word8, Word8, Word8, Word8)
NS.hostAddressToTuple Word32
host
      w :: Word32
w =
        (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
a Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
24)
          Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
16)
          Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. (Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
c Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftL` Int
8)
          Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.|. Word8 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
d
      p :: Word16
p = PortNumber -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral PortNumber
port :: Word16
  Multiaddr -> IO Multiaddr
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Multiaddr -> IO Multiaddr) -> Multiaddr -> IO Multiaddr
forall a b. (a -> b) -> a -> b
$ [Protocol] -> Multiaddr
Multiaddr [Word32 -> Protocol
IP4 Word32
w, Word16 -> Protocol
TCP Word16
p]
sockAddrToMultiaddr (NS.SockAddrInet6 PortNumber
port Word32
_ HostAddress6
host6 Word32
_) = do
  let p :: Word16
p = PortNumber -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral PortNumber
port :: Word16
      ipv6 :: IPv6
ipv6 = HostAddress6 -> IPv6
fromHostAddress6 HostAddress6
host6
      bs :: ByteString
bs = IPv6 -> ByteString
ipv6ToBytes IPv6
ipv6
  Multiaddr -> IO Multiaddr
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Multiaddr -> IO Multiaddr) -> Multiaddr -> IO Multiaddr
forall a b. (a -> b) -> a -> b
$ [Protocol] -> Multiaddr
Multiaddr [ByteString -> Protocol
IP6 ByteString
bs, Word16 -> Protocol
TCP Word16
p]
sockAddrToMultiaddr SockAddr
other =
  String -> IO Multiaddr
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO Multiaddr) -> String -> IO Multiaddr
forall a b. (a -> b) -> a -> b
$ String
"sockAddrToMultiaddr: unsupported address type: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> SockAddr -> String
forall a. Show a => a -> String
show SockAddr
other

-- | Extract the nth octet from a Word32 (0 = least significant).
octet :: Int -> Word32 -> Word8
octet :: Int -> Word32 -> Word8
octet Int
n Word32
w = Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32
w Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8)) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
0xff)

-- | Render an IPv4 Word32 to dotted-decimal string.
renderIPv4 :: Word32 -> String
renderIPv4 :: Word32 -> String
renderIPv4 Word32
w =
  Word8 -> String
forall a. Show a => a -> String
show (Int -> Word32 -> Word8
octet Int
3 Word32
w) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"." String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String
forall a. Show a => a -> String
show (Int -> Word32 -> Word8
octet Int
2 Word32
w)
    String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"." String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String
forall a. Show a => a -> String
show (Int -> Word32 -> Word8
octet Int
1 Word32
w) String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"." String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Word8 -> String
forall a. Show a => a -> String
show (Int -> Word32 -> Word8
octet Int
0 Word32
w)

-- | Render a 16-byte IPv6 ByteString to text string.
renderIPv6 :: BS.ByteString -> String
renderIPv6 :: ByteString -> String
renderIPv6 ByteString
bs = IPv6 -> String
forall a. Show a => a -> String
show (ByteString -> IPv6
bytesToIPv6 ByteString
bs)

-- | Convert 16-byte ByteString to IPv6 address.
bytesToIPv6 :: BS.ByteString -> IPv6
bytesToIPv6 :: ByteString -> IPv6
bytesToIPv6 ByteString
bs =
  let readW32 :: Int -> a
readW32 Int
off =
        (Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs Int
off) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
24)
          a -> a -> a
forall a. Bits a => a -> a -> a
.|. (Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
16)
          a -> a -> a
forall a. Bits a => a -> a -> a
.|. (Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
8)
          a -> a -> a
forall a. Bits a => a -> a -> a
.|. Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
BS.index ByteString
bs (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3))
  in HostAddress6 -> IPv6
fromHostAddress6 (Int -> Word32
forall {a}. (Bits a, Num a) => Int -> a
readW32 Int
0, Int -> Word32
forall {a}. (Bits a, Num a) => Int -> a
readW32 Int
4, Int -> Word32
forall {a}. (Bits a, Num a) => Int -> a
readW32 Int
8, Int -> Word32
forall {a}. (Bits a, Num a) => Int -> a
readW32 Int
12)

-- | Convert IPv6 address to 16-byte ByteString.
ipv6ToBytes :: IPv6 -> BS.ByteString
ipv6ToBytes :: IPv6 -> ByteString
ipv6ToBytes IPv6
ipv6 =
  let (Word32
w0, Word32
w1, Word32
w2, Word32
w3) = IPv6 -> HostAddress6
toHostAddress6 IPv6
ipv6
      word32ToBytes :: a -> [a]
word32ToBytes a
w =
        [ a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
w a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
24)
        , a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
w a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
16)
        , a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a
w a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
8)
        , a -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
w
        ]
  in [Word8] -> ByteString
BS.pack ([Word8] -> ByteString) -> [Word8] -> ByteString
forall a b. (a -> b) -> a -> b
$ Word32 -> [Word8]
forall {a} {a}. (Integral a, Bits a, Num a) => a -> [a]
word32ToBytes Word32
w0 [Word8] -> [Word8] -> [Word8]
forall a. Semigroup a => a -> a -> a
<> Word32 -> [Word8]
forall {a} {a}. (Integral a, Bits a, Num a) => a -> [a]
word32ToBytes Word32
w1 [Word8] -> [Word8] -> [Word8]
forall a. Semigroup a => a -> a -> a
<> Word32 -> [Word8]
forall {a} {a}. (Integral a, Bits a, Num a) => a -> [a]
word32ToBytes Word32
w2 [Word8] -> [Word8] -> [Word8]
forall a. Semigroup a => a -> a -> a
<> Word32 -> [Word8]
forall {a} {a}. (Integral a, Bits a, Num a) => a -> [a]
word32ToBytes Word32
w3