{-# LANGUAGE OverloadedStrings #-}
module Network.Wai.Handler.Warp.WithApplication (
  withApplication,
  withApplicationSettings,
  testWithApplication,
  testWithApplicationSettings,
  openFreePort,
  withFreePort,
) where

import           Control.Concurrent
import qualified UnliftIO
import           UnliftIO.Async
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

-- | 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 :: 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 :: 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 :: IO ()
settingsBeforeMainLoop
              = Waiter () -> () -> IO ()
forall a. Waiter a -> a -> IO ()
notify Waiter ()
started () IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Settings -> IO ()
settingsBeforeMainLoop Settings
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 (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 (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 :: 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 :: 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 (m :: * -> *) a. Monad m => a -> m a
return Application
wrappedApp) Port -> IO a
action

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

mkWaiter :: IO (Waiter a)
mkWaiter :: IO (Waiter a)
mkWaiter = do
  MVar a
mvar <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
  Waiter a -> IO (Waiter a)
forall (m :: * -> *) a. Monad m => a -> m a
return Waiter :: forall a. (a -> IO ()) -> IO a -> Waiter a
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 :: ((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)