{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Network.Test.Common
  ( -- * Client server configuration
    ClientServer(..)
  , setClientAction
  , setServerAction
  , tcp
  , unix
  , unixWithUnlink
  , udp
  , withPort
  -- * Run a ClientServer configuration
  , test
  , tcpTest
  , tcpTest6
  , udpTest
  , udpTest6
  -- * Common constants
  , serverAddr
  , serverAddr6
  , unixAddr
  , testMsg
  , lazyTestMsg
  ) where

#if __GLASGOW_HASKELL__ < 709
import Control.Applicative ((<$>))
#endif
import Control.Concurrent (ThreadId, forkIO, myThreadId)
import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar, readMVar)
import qualified Control.Exception as E
import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as L
import Network.Socket
import System.Directory
import System.Timeout (timeout)
import Test.Hspec

serverAddr :: String
serverAddr = "127.0.0.1"

serverAddr6 :: String
serverAddr6 = "::1"

testMsg :: ByteString
testMsg = "This is a test message."

lazyTestMsg :: L.ByteString
lazyTestMsg = L.fromStrict "This is a test message."

unixAddr :: String
unixAddr = "/tmp/network-test"

-- | Establish a connection between client and server and then run
-- 'clientAct' and 'serverAct', in different threads.  Both actions
-- get passed a connected 'Socket', used for communicating between
-- client and server.  'unixTest' makes sure that the 'Socket' is
-- closed after the actions have run.
unixWithUnlink :: String -> ((Socket, SockAddr) -> IO b) -> (ClientServer Socket b)
unixWithUnlink address = unix address unlink
  where
    unlink file = do
        exist <- doesFileExist file
        when exist $ removeFile file

unix
    :: String -- ^ address
    -> (String -> IO ()) -- ^ clean up action
    -> ((Socket, SockAddr) -> IO b) -- ^ server action
    -> (ClientServer Socket b)
unix address cleanupAct serverAct = defaultClientServer
    { clientSetup = do
        sock <- socket AF_UNIX Stream defaultProtocol
        connect sock (SockAddrUnix address)
        return sock
    , serverSetup = do
        sock <- socket AF_UNIX Stream defaultProtocol
        cleanupAct address -- just in case
        bind sock (SockAddrUnix address)
        listen sock 1
        return sock
    , serverAction = \sock ->
        E.bracket (accept sock) (killClientSock . fst) serverAct
    }
  where
    killClientSock sock = do
        shutdown sock ShutdownBoth
        close sock
        cleanupAct address

-- | Establish a connection between client and server and then run
-- 'clientAct' and 'serverAct', in different threads.  Both actions
-- get passed a connected 'Socket', used for communicating between
-- client and server.  'tcpTest' makes sure that the 'Socket' is
-- closed after the actions have run.
tcpTest :: (Socket -> IO a) -> (Socket -> IO b) -> IO ()
tcpTest client server = withPort $ test . setClientAction client . tcp serverAddr server

tcpTest6 :: (Socket -> IO a) -> (Socket -> IO b) -> IO ()
tcpTest6 client server = withPort $ test . setClientAction client . tcp serverAddr6 server

tcp :: HostName -> (Socket -> IO b) -> MVar PortNumber -> ClientServer Socket ()
tcp serverAddress serverAct portVar = defaultClientServer
    { clientSetup = do
        serverPort <- readMVar portVar
        addr <- resolveClient Stream serverAddress serverPort
        sock <- socketWithAddrInfo addr
#if !defined(mingw32_HOST_OS)
        withFdSocket sock $ \fd -> do
          getNonBlock fd `shouldReturn` True
          getCloseOnExec fd `shouldReturn` False
#endif
        connect sock $ addrAddress addr
        return sock
    , serverSetup = do
        addr <- resolveServer Stream serverAddress
        sock <- socketWithAddrInfo addr
        withFdSocket sock $ \fd -> do
#if !defined(mingw32_HOST_OS)
          getNonBlock fd `shouldReturn` True
          getCloseOnExec fd `shouldReturn` False
#endif
          setSocketOption sock ReuseAddr 1
          setCloseOnExecIfNeeded fd
#if !defined(mingw32_HOST_OS)
          getCloseOnExec fd `shouldReturn` True
#endif
        bind sock $ addrAddress addr
        listen sock 1
        serverPort <- socketPort sock
        putMVar portVar serverPort
        return sock
    , serverAction = \sock -> do
        (clientSock, _) <- accept sock
#if !defined(mingw32_HOST_OS)
        withFdSocket sock $ \fd -> do
          getNonBlock fd `shouldReturn` True
          getCloseOnExec fd `shouldReturn` True
#endif
        _ <- serverAct clientSock
        close clientSock
    }

-- | Create an unconnected 'Socket' for sending UDP and receiving
-- datagrams and then run 'clientAct' and 'serverAct'.
udpTest :: (Socket -> SockAddr -> IO a) -> (Socket -> IO b) -> IO ()
udpTest client server =
    withPort $ test . setServerAction server . udp serverAddr client

udpTest6 :: (Socket -> SockAddr -> IO a) -> (Socket -> IO b) -> IO ()
udpTest6 client server =
    withPort $ test . setServerAction server . udp serverAddr6 client

udp
    :: HostName
    -> (Socket -> SockAddr -> IO a)
    -> MVar PortNumber
    -> ClientServer a Socket
udp serverAddress clientAct portVar = defaultClientServer
    { clientSetup = do
        addr <- resolveClient Datagram serverAddress 8000 -- dummy port
        socketWithAddrInfo addr
    , clientAction = \sock -> do
        serverPort <- readMVar portVar
        addr <- resolveClient Datagram serverAddress serverPort
        clientAct sock $ addrAddress addr
    , serverSetup = do
        addr <- resolveServer Datagram serverAddress
        sock <- socketWithAddrInfo addr
        setSocketOption sock ReuseAddr 1
        bind sock $ addrAddress addr
        serverPort <- socketPort sock
        putMVar portVar serverPort
        return sock
    }

data ClientServer a b
    = ClientServer
    { clientSetup :: IO Socket
    , clientAction :: Socket -> IO a
    , serverSetup :: IO Socket
    , serverAction :: Socket -> IO b
    }

setClientAction
    :: (Socket -> IO b)
    -> ClientServer a c
    -> ClientServer b c
setClientAction f c = c { clientAction = f }

setServerAction
    :: (Socket -> IO c)
    -> ClientServer a b
    -> ClientServer a c
setServerAction f c = c { serverAction = f }

defaultClientServer :: ClientServer Socket Socket
defaultClientServer = ClientServer
    { clientSetup =
        E.throwIO $ userError "no client setup defined"
    , clientAction = return
    , serverSetup = E.throwIO $ userError "no server setup defined"
    , serverAction = return
    }

-- | Run a client/server pair and synchronize them so that the server
-- is started before the client and the specified server action is
-- finished before the client closes the 'Socket'.
test :: ClientServer a b -> IO ()
test conf = do
    tid <- myThreadId
    barrier <- newEmptyMVar
    _ <- forkIO $ server tid barrier
    client tid barrier
  where
    server tid barrier =
        bracketWithReraise tid (serverSetup conf) close $ \sock -> do
        serverReady
        Just _ <- timeout 1000000 $ (serverAction conf) sock
        putMVar barrier ()
      where
        -- | Signal to the client that it can proceed.
        serverReady = putMVar barrier ()
    client tid barrier = do
        takeMVar barrier
        -- Transfer exceptions to the main thread.
        bracketWithReraise tid (clientSetup conf) close $ \res -> do
            Just _ <- timeout 1000000 $ (clientAction conf) res
            takeMVar barrier

withPort :: (MVar PortNumber -> IO a) -> IO a
withPort f = f =<< newEmptyMVar

-- | Like 'bracket' but catches and reraises the exception in another
-- thread, specified by the first argument.
bracketWithReraise :: ThreadId -> IO a -> (a -> IO b) -> (a -> IO ()) -> IO ()
bracketWithReraise tid setup teardown thing =
    E.bracket setup teardown thing
    `E.catch` \ (e :: E.SomeException) -> E.throwTo tid e

resolveClient :: SocketType -> HostName -> PortNumber -> IO AddrInfo
resolveClient socketType host port =
    head <$> getAddrInfo (Just hints) (Just host) (Just $ show port)
  where
    hints = defaultHints {
        addrSocketType = socketType
      , addrFlags = [AI_NUMERICHOST]
      }

resolveServer :: SocketType -> HostName -> IO AddrInfo
resolveServer socketType host =
    head <$> getAddrInfo (Just hints) (Just host) Nothing
  where
    hints = defaultHints {
        addrSocketType = socketType
      , addrFlags = [AI_PASSIVE]
      }

socketWithAddrInfo :: AddrInfo -> IO Socket
socketWithAddrInfo addr = socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)