{-# LANGUAGE OverloadedStrings #-}

module Network.Wai.Handler.Warp.WithApplication (
    withApplication,
    withApplicationSettings,
    testWithApplication,
    testWithApplicationSettings,
    openFreePort,
    withFreePort,
) where

import Control.Concurrent
import Control.Monad (when)
import Data.Streaming.Network (bindRandomPortTCP)
import Network.Socket
import Network.Wai
import Network.Wai.Handler.Warp.Run
import Network.Wai.Handler.Warp.Settings
import Network.Wai.Handler.Warp.Types
import qualified UnliftIO
import UnliftIO.Async

-- | Runs the given 'Application' on a free port. Passes the port to the given
-- operation and executes it, while the 'Application' is running. Shuts down the
-- server before returning.
--
-- @since 3.2.4
withApplication :: IO Application -> (Port -> IO a) -> IO a
withApplication :: forall a. IO Application -> (Port -> IO a) -> IO a
withApplication = Settings -> IO Application -> (Port -> IO a) -> IO a
forall a. Settings -> IO Application -> (Port -> IO a) -> IO a
withApplicationSettings Settings
defaultSettings

-- | 'withApplication' with given 'Settings'. This will ignore the port value
-- set by 'setPort' in 'Settings'.
--
-- @since 3.2.7
withApplicationSettings :: Settings -> IO Application -> (Port -> IO a) -> IO a
withApplicationSettings :: forall a. Settings -> IO Application -> (Port -> IO a) -> IO a
withApplicationSettings Settings
settings' IO Application
mkApp Port -> IO a
action = do
    Application
app <- IO Application
mkApp
    ((Port, Socket) -> IO a) -> IO a
forall a. ((Port, Socket) -> IO a) -> IO a
withFreePort (((Port, Socket) -> IO a) -> IO a)
-> ((Port, Socket) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \(Port
port, Socket
sock) -> do
        Waiter ()
started <- IO (Waiter ())
forall a. IO (Waiter a)
mkWaiter
        let settings :: Settings
settings =
                Settings
settings'
                    { settingsBeforeMainLoop =
                        notify started () >> settingsBeforeMainLoop settings'
                    }
        Either () a
result <-
            IO () -> IO a -> IO (Either () a)
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> m b -> m (Either a b)
race
                (Settings -> Socket -> Application -> IO ()
runSettingsSocket Settings
settings Socket
sock Application
app)
                (Waiter () -> IO ()
forall a. Waiter a -> IO a
waitFor Waiter ()
started IO () -> IO a -> IO a
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Port -> IO a
action Port
port)
        case Either () a
result of
            Left () -> String -> IO a
forall (m :: * -> *) a. (MonadIO m, HasCallStack) => String -> m a
UnliftIO.throwString String
"Unexpected: runSettingsSocket exited"
            Right a
x -> a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x

-- | Same as 'withApplication' but with different exception handling: If the
-- given 'Application' throws an exception, 'testWithApplication' will re-throw
-- the exception to the calling thread, possibly interrupting the execution of
-- the given operation.
--
-- This is handy for running tests against an 'Application' over a real network
-- port. When running tests, it's useful to let exceptions thrown by your
-- 'Application' propagate to the main thread of the test-suite.
--
-- __The exception handling makes this function unsuitable for use in production.__
-- Use 'withApplication' instead.
--
-- @since 3.2.4
testWithApplication :: IO Application -> (Port -> IO a) -> IO a
testWithApplication :: forall a. IO Application -> (Port -> IO a) -> IO a
testWithApplication = Settings -> IO Application -> (Port -> IO a) -> IO a
forall a. Settings -> IO Application -> (Port -> IO a) -> IO a
testWithApplicationSettings Settings
defaultSettings

-- | 'testWithApplication' with given 'Settings'.
--
-- @since 3.2.7
testWithApplicationSettings
    :: Settings -> IO Application -> (Port -> IO a) -> IO a
testWithApplicationSettings :: forall a. Settings -> IO Application -> (Port -> IO a) -> IO a
testWithApplicationSettings Settings
settings IO Application
mkApp Port -> IO a
action = do
    ThreadId
callingThread <- IO ThreadId
myThreadId
    Application
app <- IO Application
mkApp
    let wrappedApp :: Application
wrappedApp Request
request Response -> IO ResponseReceived
respond =
            Application
app Request
request Response -> IO ResponseReceived
respond IO ResponseReceived
-> (SomeException -> IO ResponseReceived) -> IO ResponseReceived
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> (SomeException -> m a) -> m a
`UnliftIO.catchAny` \SomeException
e -> do
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when
                    (SomeException -> Bool
defaultShouldDisplayException SomeException
e)
                    (ThreadId -> SomeException -> IO ()
forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
callingThread SomeException
e)
                SomeException -> IO ResponseReceived
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
UnliftIO.throwIO SomeException
e
    Settings -> IO Application -> (Port -> IO a) -> IO a
forall a. Settings -> IO Application -> (Port -> IO a) -> IO a
withApplicationSettings Settings
settings (Application -> IO Application
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Application
wrappedApp) Port -> IO a
action

data Waiter a = Waiter
    { forall a. Waiter a -> a -> IO ()
notify :: a -> IO ()
    , forall a. Waiter a -> IO a
waitFor :: IO a
    }

mkWaiter :: IO (Waiter a)
mkWaiter :: forall a. IO (Waiter a)
mkWaiter = do
    MVar a
mvar <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
    Waiter a -> IO (Waiter a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
        Waiter
            { notify :: a -> IO ()
notify = MVar a -> a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar a
mvar
            , waitFor :: IO a
waitFor = MVar a -> IO a
forall a. MVar a -> IO a
readMVar MVar a
mvar
            }

-- | Opens a socket on a free port and returns both port and socket.
--
-- @since 3.2.4
openFreePort :: IO (Port, Socket)
openFreePort :: IO (Port, Socket)
openFreePort = HostPreference -> IO (Port, Socket)
bindRandomPortTCP HostPreference
"127.0.0.1"

-- | Like 'openFreePort' but closes the socket before exiting.
withFreePort :: ((Port, Socket) -> IO a) -> IO a
withFreePort :: forall a. ((Port, Socket) -> IO a) -> IO a
withFreePort = IO (Port, Socket)
-> ((Port, Socket) -> IO ()) -> ((Port, Socket) -> IO a) -> IO a
forall (m :: * -> *) a b c.
MonadUnliftIO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
UnliftIO.bracket IO (Port, Socket)
openFreePort (Socket -> IO ()
close (Socket -> IO ())
-> ((Port, Socket) -> Socket) -> (Port, Socket) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Port, Socket) -> Socket
forall a b. (a, b) -> b
snd)