{-# LANGUAGE CPP #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} module Network.Test.Common ( -- * Client server configuration ClientServer(..) , setClientAction , setServerAction , defaultClientServer , 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 qualified Network.Socket.Address 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)