@@ -10,7 +10,7 @@ import Control.Arrow (first)
1010import Control.Concurrent (threadDelay, forkIOWithUnmask)
1111import qualified Control.Concurrent as Conc (yield)
1212import Control.Exception as E
13−import Control.Monad (forever, when, unless, void)
13+import Control.Monad (when, unless, void)
1414import Data.ByteString (ByteString)
1515import qualified Data.ByteString as S
1616import Data.IORef (IORef, newIORef, readIORef, writeIORef)
@@ -20,6 +20,7 @@ import Network.Socket (accept, withSocketsDo, SockAddr)
2020import qualified Network.Socket.ByteString as Sock
2121import Network.Wai
2222import Network.Wai.Handler.Warp.Buffer
23+import Network.Wai.Handler.Warp.Counter
2324import qualified Network.Wai.Handler.Warp.Date as D
2425import qualified Network.Wai.Handler.Warp.FdCache as F
2526import Network.Wai.Handler.Warp.Header
@@ -31,6 +32,7 @@ import Network.Wai.Handler.Warp.Settings
3132import qualified Network.Wai.Handler.Warp.Timeout as T
3233import Network.Wai.Handler.Warp.Types
3334import Network.Wai.Internal (ResponseReceived (ResponseReceived))
35+import System.IO.Error (isFullErrorType, ioeGetErrorType)
3436
3537#if WINDOWS
3638import Network.Wai.Handler.Warp.Windows
@@ -82,8 +84,12 @@ runSettings set app = withSocketsDo $
8284--
8385-- Note that the 'settingsPort' will still be passed to 'Application's via the
8486-- 'serverPort' record.
87+--
88+-- When the listen socket in the second argument is closed, all live
89+-- connections are gracefully shut-downed.
8590runSettingsSocket :: Settings -> Socket -> Application -> IO ()
86−runSettingsSocket set socket app =
91+runSettingsSocket set socket app = do
92+ settingsInstallShutdownHandler set closeListenSocket
8793 runSettingsConnection set getConn app
8894 where
8995 getConn = do
@@ -96,6 +102,8 @@ runSettingsSocket set socket app =
96102 conn <- socketConnection s
97103 return (conn, sa)
98104
105+ closeListenSocket = sClose socket
106+
99107-- | Allows you to provide a function which will return a 'Connection'. In
100108-- cases where creating the @Connection@ can be expensive, this allows the
101109-- expensive computations to be performed in a separate thread instead of the
@@ -115,6 +123,8 @@ runSettingsConnectionMaker x y =
115123 where
116124 go = fmap (first (fmap (, False)))
117125
126+----------------------------------------------------------------
127+
118128-- | Allows you to provide a function which will return a function
119129-- which will return 'Connection'.
120130--
@@ -122,95 +132,137 @@ runSettingsConnectionMaker x y =
122132runSettingsConnectionMakerSecure :: Settings -> IO (IO (Connection, Bool), SockAddr) -> Application -> IO ()
123133runSettingsConnectionMakerSecure set getConnMaker app = do
124134 settingsBeforeMainLoop set
135+ counter <- newCounter
125136
126− -- Note that there is a thorough discussion of the exception safety of the
127− -- following code at: https://github.com/yesodweb/wai/issues/146
128− --
129− -- We need to make sure of two things:
130− --
131− -- 1. Asynchronous exceptions are not blocked entirely in the main loop.
132− -- Doing so would make it impossible to kill the Warp thread.
133− --
134− -- 2. Once a connection maker is received via getConnLoop, the connection
135− -- is guaranteed to be closed, even in the presence of async exceptions.
136− --
137− -- Our approach is explained in the comments below.
138−
139− -- First mask all exceptions in the main loop. This is necessary to ensure
140− -- that no async exception is throw between the call to getConnLoop and the
141− -- registering of connClose.
142− D.withDateCache $ \dc -> do
143− F.withFdCache (settingsFdCacheDuration set * 1000000) $ \fc -> do
144− withTimeoutManager $ \tm -> mask_ . forever $ do
137+ D.withDateCache $ \dc ->
138+ F.withFdCache fdCacheDurationInSeconds $ \fc ->
139+ withTimeoutManager $ \tm ->
140+ acceptConnection set getConnMaker app dc fc tm counter
141+ where
142+ fdCacheDurationInSeconds = settingsFdCacheDuration set * 1000000
143+ withTimeoutManager f = case settingsManager set of
144+ Just tm -> f tm
145+ Nothing -> bracket
146+ (T.initialize $ settingsTimeout set * 1000000)
147+ T.stopManager
148+ f
149+
150+onE :: Settings -> Maybe Request -> SomeException -> IO ()
151+onE set mreq e = case fromException e of
152+ Just (NotEnoughLines []) -> return ()
153+ _ -> settingsOnException set mreq e
154+
155+-- Note that there is a thorough discussion of the exception safety of the
156+-- following code at: https://github.com/yesodweb/wai/issues/146
157+--
158+-- We need to make sure of two things:
159+--
160+-- 1. Asynchronous exceptions are not blocked entirely in the main loop.
161+-- Doing so would make it impossible to kill the Warp thread.
162+--
163+-- 2. Once a connection maker is received via acceptNewConnection, the
164+-- connection is guaranteed to be closed, even in the presence of
165+-- async exceptions.
166+--
167+-- Our approach is explained in the comments below.
168+acceptConnection :: Settings
169+ -> IO (IO (Connection, Bool), SockAddr)
170+ -> Application
171+ -> D.DateCache
172+ -> Maybe F.MutableFdCache
173+ -> T.Manager
174+ -> Counter
175+ -> IO ()
176+acceptConnection set getConnMaker app dc fc tm counter = do
177+ -- First mask all exceptions in acceptLoop. This is necessary to
178+ -- ensure that no async exception is throw between the call to
179+ -- acceptNewConnection and the registering of connClose.
180+ void $ mask_ $ acceptLoop
181+ gracefulShutdown counter
182+ where
183+ acceptLoop = do
145184 -- Allow async exceptions before receiving the next connection maker.
146185 allowInterrupt
147186
148− -- getConnLoop will try to receive the next incoming request. It
149− -- returns a /connection maker/, not a connection, since in some
150− -- circumstances creating a working connection from a raw socket may be
151− -- an expensive operation, and this expensive work should not be
152− -- performed in the main event loop. An example of something expensive
153− -- would be TLS negotiation.
154− (mkConn, addr) <- getConnLoop
155−
156− -- Fork a new worker thread for this connection maker, and ask for a
157− -- function to unmask (i.e., allow async exceptions to be thrown).
158− --
159− -- GHC 7.8 cannot infer the type of "void . forkIOWithUnmask"
160− void $ forkIOWithUnmask $ \unmask ->
161− -- Run the connection maker to get a new connection, and ensure
162− -- that the connection is closed. If the mkConn call throws an
163− -- exception, we will leak the connection. If the mkConn call is
164− -- vulnerable to attacks (e.g., Slowloris), we do nothing to
165− -- protect the server. It is therefore vital that mkConn is well
166− -- vetted.
167− --
168− -- We grab the connection before registering timeouts since the
169− -- timeouts will be useless during connection creation, due to the
170− -- fact that async exceptions are still masked.
171− bracket mkConn (connClose . fst) $ \(conn', isSecure') ->
172−
173− -- We need to register a timeout handler for this thread, and
174− -- cancel that handler as soon as we exit.
175− bracket (T.registerKillThread tm) T.cancel $ \th ->
176− let ii = InternalInfo th fc dc
177− conn = setSendFile conn' fc
178− -- We now have fully registered a connection close handler
179− -- in the case of all exceptions, so it is safe to one
180− -- again allow async exceptions.
181− in unmask .
182− -- Call the user-supplied on exception code if any
183− -- exceptions are thrown.
184− handle (onE Nothing) .
185−
186− -- Call the user-supplied code for connection open and close events
187− bracket (onOpen addr) (const $ onClose addr) $ \goingon ->
188−
189− -- Actually serve this connection.
190− -- onnClose above ensures the termination of the connection.
191− when goingon $ serveConnection conn ii addr isSecure' set app
187+ -- acceptNewConnection will try to receive the next incoming
188+ -- request. It returns a /connection maker/, not a connection,
189+ -- since in some circumstances creating a working connection
190+ -- from a raw socket may be an expensive operation, and this
191+ -- expensive work should not be performed in the main event
192+ -- loop. An example of something expensive would be TLS
193+ -- negotiation.
194+ mx <- acceptNewConnection
195+ case mx of
196+ Nothing -> return ()
197+ Just (mkConn, addr) -> do
198+ fork set mkConn addr app dc fc tm counter
199+ acceptLoop
200+
201+ acceptNewConnection = do
202+ ex <- try getConnMaker
203+ case ex of
204+ Right x -> return $ Just x
205+ Left e -> do
206+ onE set Nothing $ toException e
207+ if isFullErrorType (ioeGetErrorType e) then do
208+ -- "resource exhausted (Too many open files)" may
209+ -- happen by accept(). Wait a second hoping that
210+ -- resource will be available.
211+ threadDelay 1000000
212+ acceptNewConnection
213+ else
214+ -- Assuming the listen socket is closed.
215+ return Nothing
216+
217+-- Fork a new worker thread for this connection maker, and ask for a
218+-- function to unmask (i.e., allow async exceptions to be thrown).
219+fork :: Settings
220+ -> IO (Connection, Bool)
221+ -> SockAddr
222+ -> Application
223+ -> D.DateCache
224+ -> Maybe F.MutableFdCache
225+ -> T.Manager
226+ -> Counter
227+ -> IO ()
228+fork set mkConn addr app dc fc tm counter = void $ forkIOWithUnmask $ \unmask ->
229+ -- Run the connection maker to get a new connection, and ensure
230+ -- that the connection is closed. If the mkConn call throws an
231+ -- exception, we will leak the connection. If the mkConn call is
232+ -- vulnerable to attacks (e.g., Slowloris), we do nothing to
233+ -- protect the server. It is therefore vital that mkConn is well
234+ -- vetted.
235+ --
236+ -- We grab the connection before registering timeouts since the
237+ -- timeouts will be useless during connection creation, due to the
238+ -- fact that async exceptions are still masked.
239+ bracket mkConn closeConn $ \(conn0, isSecure') ->
240+
241+ -- We need to register a timeout handler for this thread, and
242+ -- cancel that handler as soon as we exit.
243+ bracket (T.registerKillThread tm) T.cancel $ \th ->
244+
245+ let ii = InternalInfo th fc dc
246+ conn = setSendFile conn0 fc
247+ -- We now have fully registered a connection close handler
248+ -- in the case of all exceptions, so it is safe to one
249+ -- again allow async exceptions.
250+ in unmask .
251+ -- Call the user-supplied on exception code if any
252+ -- exceptions are thrown.
253+ handle (onE set Nothing) .
254+
255+ -- Call the user-supplied code for connection open and close events
256+ bracket (onOpen addr) (onClose addr) $ \goingon ->
257+
258+ -- Actually serve this connection.
259+ -- onnClose above ensures the termination of the connection.
260+ when goingon $ serveConnection conn ii addr isSecure' set app
192261 where
193− -- FIXME: only IOEception is caught. What about other exceptions?
194− getConnLoop = getConnMaker `E.catch` \(e :: IOException) -> do
195− onE Nothing (toException e)
196− -- "resource exhausted (Too many open files)" may happen by accept().
197− -- Wait a second hoping that resource will be available.
198− threadDelay 1000000
199− getConnLoop
200− onE mreq e =
201− case fromException e of
202− Just (NotEnoughLines []) -> return ()
203− _ -> settingsOnException set mreq e
204− onOpen = settingsOnOpen set
205− onClose = settingsOnClose set
206−
207− withTimeoutManager f =
208− case settingsManager set of
209− Nothing -> bracket
210− (T.initialize $ settingsTimeout set * 1000000)
211− T.stopManager
212− f
213− Just tm -> f tm
262+ closeConn (conn, _isSecure) = connClose conn
263+
264+ onOpen adr = increase counter >> settingsOnOpen set adr
265+ onClose adr _ = decrease counter >> settingsOnClose set adr
214266
215267serveConnection :: Connection
216268 -> InternalInfo
@@ -299,3 +351,10 @@ setSocketCloseOnExec _ = return ()
299351setSocketCloseOnExec socket =
300352 setFdOption (fromIntegral $ fdSocket socket) CloseOnExec True
301353#endif
354+
355+gracefulShutdown :: Counter -> IO ()
356+gracefulShutdown counter = do
357+ -- To avoid race condition, we just use threadDelay, not MVar.
358+ threadDelay 10000000
359+ noConnections <- isZero counter
360+ unless noConnections $ gracefulShutdown counter