{-# LANGUAGE LambdaCase #-}

module Protocol.RT.Network where

-- external, impure
import qualified Network.Socket        as N

import           Network.Socket        hiding (PortNumber, SockAddr (..))

-- external, pure
import qualified Data.ByteString.Char8 as BSC

import           Data.Bits
import           Data.Word
import           Protocol.Base
import           Safe.Numeric


toNPort :: PortNumber -> N.PortNumber
toNPort = fromIntegral

fromNPort :: N.PortNumber -> PortNumber
fromNPort = fromIntegral

-- HostAddress is in host-order so we need to switch endian
-- the architecture-specific logic is performed in hostAddressToTuple
fromNIpv4 :: HostAddress -> IpAddr4
fromNIpv4 h =
  let (a, b, c, d) = hostAddressToTuple h
  in  ex d .|. (ex c `shiftL` 8) .|. (ex b `shiftL` 16) .|. (ex a `shiftL` 24)

-- HostAddress is in host-order so we need to switch endian
-- the architecture-specific logic is performed in tupleToHostAddress
toNIpv4 :: IpAddr4 -> HostAddress
toNIpv4 p =
  let a = ctX (p `shiftR` 24)
      b = ctW (p `shiftR` 16)
      c = ctW (p `shiftR` 8)
      d = ctW p
  in  tupleToHostAddress (a, b, c, d)

w64to32 :: Word64 -> (Word32, Word32)
w64to32 w = (ctX (w `shiftR` 32), ctW w)

w32to64 :: Word32 -> Word32 -> Word64
w32to64 a b = (ex a `shiftL` 32) .|. ex b

-- HostAddress6 is already in network-order, so we don't need to switch endian
toNIpv6 :: IpAddr6 -> HostAddress6
toNIpv6 (Word128 hi lo) =
  let ((a, b), (c, d)) = (w64to32 hi, w64to32 lo) in (a, b, c, d)

-- HostAddress6 is already in network-order, so we don't need to switch endian
fromNIpv6 :: HostAddress6 -> IpAddr6
fromNIpv6 (a, b, c, d) = Word128 (w32to64 a b) (w32to64 c d)

fromNAddr :: N.SockAddr -> SockAddr
fromNAddr = \case
  N.SockAddrInet p h -> SockAddrInet4 (SockAddr4 (fromNIpv4 h) (fromNPort p))
  N.SockAddrInet6 p f h s ->
    SockAddrInet6 (SockAddr6 (fromNIpv6 h) (fromNPort p) f s)
  N.SockAddrUnix p -> SockAddrUnix (BSC.pack p)

toNAddr :: SockAddr -> N.SockAddr
toNAddr = \case
  SockAddrInet4 (SockAddr4 h p) -> N.SockAddrInet (toNPort p) (toNIpv4 h)
  SockAddrInet6 (SockAddr6 h p f s) ->
    N.SockAddrInet6 (toNPort p) f (toNIpv6 h) s
  SockAddrUnix p -> N.SockAddrUnix (BSC.unpack p)

sockAddrFamily :: N.SockAddr -> Family
sockAddrFamily addr = case addr of
  N.SockAddrInet{}  -> AF_INET
  N.SockAddrInet6{} -> AF_INET6
  N.SockAddrUnix _  -> AF_UNIX

setSockOptsForQuic :: Socket -> IO ()
setSockOptsForQuic sock = getSocketName sock >>= \case
  -- TODO: quinn (not quinn-proto) also sets MTU_DISCOVER for some reason
  -- anyway, it's not available in Network.Socket, ignore it for now & hope nothing breaks
  N.SockAddrInet{} -> do
    setSocketOption sock RecvIPv4TOS     1
    setSocketOption sock RecvIPv4PktInfo 1
  N.SockAddrInet6{} -> do
    setSocketOption sock RecvIPv6TClass  1
    setSocketOption sock RecvIPv6PktInfo 1
  N.SockAddrUnix _ -> pure ()

-- see also https://blog.powerdns.com/2012/10/08/on-binding-datagram-udp-sockets-to-the-any-addresses/

getDstIp :: [Cmsg] -> N.SockAddr -> Maybe IpAddr
getDstIp cmsg = \case
  N.SockAddrInet{} -> case lookupCmsg CmsgIdIPv4PktInfo cmsg of
    Nothing -> Nothing
    Just cm -> case decodeCmsg cm of
      Nothing                         -> Nothing
      Just (IPv4PktInfo _ _ ipi_addr) -> Just $ IpAddr4 (fromNIpv4 ipi_addr)
  N.SockAddrInet6{} -> case lookupCmsg CmsgIdIPv6PktInfo cmsg of
    Nothing -> Nothing
    Just cm -> case decodeCmsg cm of
      Nothing                       -> Nothing
      Just (IPv6PktInfo _ dst_addr) -> Just $ IpAddr6 $ fromNIpv6 dst_addr
  N.SockAddrUnix _ -> Nothing

setSrcIp :: N.SockAddr -> Maybe IpAddr -> [Cmsg]
setSrcIp addr = \case
  Nothing    -> []
  Just srcIp -> case addr of
    N.SockAddrInet{} -> case srcIp of
      IpAddr4 src -> [encodeCmsg $ IPv4PktInfo 0 0 (toNIpv4 src)]
      _           -> error "setSrcIp: non-IPv4 srcIp for IPv4 socket"
    N.SockAddrInet6{} -> case srcIp of
      IpAddr6 src -> [encodeCmsg $ IPv6PktInfo 0 $ toNIpv6 src]
      _           -> error "setSrcIp: non-IPv6 srcIp for IPv6 socket"
    N.SockAddrUnix _ -> []

getEcn :: [Cmsg] -> N.SockAddr -> Maybe EcnCodepoint
getEcn cmsg = \case
  N.SockAddrInet{} -> case lookupCmsg CmsgIdIPv4TOS cmsg of
    Nothing -> Nothing
    Just cm -> case decodeCmsg cm of
      Nothing          -> Nothing
      Just (IPv4TOS c) -> ecnFromBits c
  N.SockAddrInet6{} -> case lookupCmsg CmsgIdIPv6TClass cmsg of
    Nothing -> Nothing
    Just cm -> case decodeCmsg cm of
      Nothing             -> Nothing
      Just (IPv6TClass c) -> ecnFromBits c
  N.SockAddrUnix _ -> Nothing

setEcn :: N.SockAddr -> Maybe EcnCodepoint -> [Cmsg]
setEcn addr = \case
  Nothing  -> []
  Just ecn -> do
    let ecn' = fromEnum ecn
    case addr of
      N.SockAddrInet{}  -> [encodeCmsg $ IPv4TOS $ fromIntegral ecn']
      N.SockAddrInet6{} -> [encodeCmsg $ IPv6TClass $ fromIntegral ecn']
      N.SockAddrUnix _  -> []

setSegmentSize :: N.SockAddr -> Maybe Word16 -> [Cmsg]
setSegmentSize addr = \case
  Nothing -> []
  Just x  -> error "not supported yet"
    -- at the time of writing:
    -- quinn-proto never sets this
    -- quinn only supports it for linux anyway currently
    -- Network.Socket does not support this on any platform