-- | TCP transport implementation.
--
-- Provides dial/listen for /ip4/.../tcp/... and /ip6/.../tcp/... multiaddrs.
-- Uses the network library for socket operations.
module Network.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 Network.LibP2P.Multiaddr.Multiaddr (Multiaddr (..))
import Network.LibP2P.Multiaddr.Protocol (Protocol (..))
import Network.LibP2P.MultistreamSelect.Negotiation (StreamIO (..))
import Network.LibP2P.Transport.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
  }

-- | Check if a multiaddr is a TCP address (/ip4/.../tcp/... or /ip6/.../tcp/...).
canDialTCP :: Multiaddr -> Bool
canDialTCP :: Multiaddr -> Bool
canDialTCP (Multiaddr [IP4 Word32
_, TCP Word16
_]) = Bool
True
canDialTCP (Multiaddr [IP6 ByteString
_, TCP Word16
_]) = Bool
True
canDialTCP Multiaddr
_ = Bool
False

-- | Extract (HostName, ServiceName) from a TCP multiaddr.
multiaddrToHostPort :: Multiaddr -> Either String (String, String)
multiaddrToHostPort :: Multiaddr -> Either String (String, String)
multiaddrToHostPort (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)
multiaddrToHostPort (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)
multiaddrToHostPort 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.
tcpDial :: Multiaddr -> IO RawConnection
tcpDial :: Multiaddr -> IO RawConnection
tcpDial addr :: Multiaddr
addr@(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
tcpDial addr :: Multiaddr
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
tcpDial 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.
mkRawConnection :: NS.Socket -> Multiaddr -> IO RawConnection
mkRawConnection :: Socket -> Multiaddr -> IO RawConnection
mkRawConnection Socket
sock Multiaddr
remoteAddr = do
  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