@@ -17,15 +17,12 @@ import qualified Control.Concurrent as Conc (yield)
1717import Control.Exception as E
1818import qualified Data.ByteString as S
1919import Data.Char (chr)
20−import Data.IORef (IORef, newIORef, readIORef, writeIORef)
20+import Data.IORef (IORef, newIORef, readIORef, writeIORef, atomicModifyIORef')
2121import Data.Streaming.Network (bindPortTCP)
2222import Foreign.C.Error (Errno(..), eCONNABORTED)
2323import GHC.IO.Exception (IOException(..))
2424import qualified Network.HTTP2 as H2
2525import Network.Socket (Socket, close, accept, withSocketsDo, SockAddr(SockAddrInet, SockAddrInet6), setSocketOption, SocketOption(..))
26−#if MIN_VERSION_network(3,1,1)
27−import Network.Socket (gracefulClose)
28−#endif
2926import qualified Network.Socket.ByteString as Sock
3027import Network.Wai
3128import Network.Wai.Internal (ResponseReceived (ResponseReceived))
@@ -58,38 +55,21 @@ import Network.Socket (fdSocket)
5855#endif
5956
6057-- | Creating 'Connection' for plain HTTP based on a given socket.
61−socketConnection :: Settings -> Socket -> IO Connection
62−#if MIN_VERSION_network(3,1,1)
63−socketConnection set s = do
64−#else
65−socketConnection _ s = do
66−#endif
58+socketConnection :: Socket -> IO Connection
59+socketConnection s = do
6760 bufferPool <- newBufferPool
6861 writeBuf <- allocateBuffer bufferSize
6962 let sendall = Sock.sendAll s
70− isH2 <- newIORef False -- HTTP/1.x
7163 return Connection {
7264 connSendMany = Sock.sendMany s
7365 , connSendAll = sendall
7466 , connSendFile = sendFile s writeBuf bufferSize sendall
75−#if MIN_VERSION_network(3,1,1)
76− , connClose = do
77− h2 <- readIORef isH2
78− let tm = if h2 then settingsGracefulCloseTimeout2 set
79− else settingsGracefulCloseTimeout1 set
80− if tm == 0 then
81− close s
82− else
83− gracefulClose s tm `E.catch` \(E.SomeException _) -> return ()
84−#else
8567 , connClose = close s
86−#endif
8768 , connFree = freeBuffer writeBuf
8869 , connRecv = receive s bufferPool
8970 , connRecvBuf = receiveBuf s
9071 , connWriteBuffer = writeBuf
9172 , connBufferSize = bufferSize
92− , connHTTP2 = isH2
9373 }
9474
9575-- | Run an 'Application' on the given port.
@@ -151,7 +131,7 @@ runSettingsSocket set socket app = do
151131 setSocketCloseOnExec s
152132 -- NoDelay causes an error for AF_UNIX.
153133 setSocketOption s NoDelay 1 `E.catch` \(E.SomeException _) -> return ()
154− conn <- socketConnection set s
134+ conn <- socketConnection s
155135 return (conn, sa)
156136
157137 closeListenSocket = close socket
@@ -286,7 +266,10 @@ fork :: Settings
286266fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
287267 -- Call the user-supplied on exception code if any
288268 -- exceptions are thrown.
289− handle (settingsOnException set Nothing) $
269+ handle (settingsOnException set Nothing) .
270+ -- Allocate a new IORef indicating whether the connection has been
271+ -- closed, to avoid double-freeing a connection
272+ withClosedRef $ \ref ->
290273 -- Run the connection maker to get a new connection, and ensure
291274 -- that the connection is closed. If the mkConn call throws an
292275 -- exception, we will leak the connection. If the mkConn call is
@@ -297,13 +280,21 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
297280 -- We grab the connection before registering timeouts since the
298281 -- timeouts will be useless during connection creation, due to the
299282 -- fact that async exceptions are still masked.
300− bracket mkConn cleanUp (serve unmask)
283+ bracket mkConn (cleanUp ref) (serve unmask ref)
301284 where
302− cleanUp (conn, _) = connClose conn `finally` connFree conn
285+ withClosedRef inner = newIORef False >>= inner
286+
287+ closeConn ref conn = do
288+ isClosed <- atomicModifyIORef' ref $ \x -> (True, x)
289+ unless isClosed $ connClose conn
290+
291+ cleanUp ref (conn, _) = closeConn ref conn `finally` connFree conn
303292
304293 -- We need to register a timeout handler for this thread, and
305− -- cancel that handler as soon as we exit.
306− serve unmask (conn, transport) = bracket register cancel $ \th -> do
294+ -- cancel that handler as soon as we exit. We additionally close
295+ -- the connection immediately in case the child thread catches the
296+ -- async exception or performs some long-running cleanup action.
297+ serve unmask ref (conn, transport) = bracket register cancel $ \th -> do
307298 -- We now have fully registered a connection close handler in
308299 -- the case of all exceptions, so it is safe to once again
309300 -- allow async exceptions.
@@ -315,7 +306,8 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask ->
315306 -- above ensures the connection is closed.
316307 when goingon $ serveConnection conn ii th addr transport set app
317308 where
318− register = T.registerKillThread (timeoutManager ii) (connClose conn)
309+ register = T.registerKillThread (timeoutManager ii)
310+ (closeConn ref conn)
319311 cancel = T.cancel
320312
321313 onOpen adr = increase counter >> settingsOnOpen set adr
@@ -342,19 +334,10 @@ serveConnection conn ii th origAddr transport settings app = do
342334 istatus <- newIORef False
343335 if settingsHTTP2Enabled settings && h2 then do
344336 rawRecvN <- makeReceiveN bs (connRecv conn) (connRecvBuf conn)
345− -- This thread becomes the sender in http2 library.
346− -- In the case of event source, one request comes and one
347− -- worker gets busy. But it is likely that the receiver does
348− -- not receive any data at all while the sender is sending
349− -- output data from the worker. It's not good enough to tickle
350− -- the time handler in the receiver only. So, we should tickle
351− -- the time handler in both the receiver and the sender.
352337 let recvN = wrappedRecvN th istatus (settingsSlowlorisSize settings) rawRecvN
353− sendBS x = connSendAll conn x >> T.tickle th
354338 -- fixme: origAddr
355339 checkTLS
356− setConnHTTP2 conn True
357− http2 conn transport ii origAddr settings recvN sendBS app
340+ http2 conn ii origAddr settings recvN app
358341 else do
359342 src <- mkSource (wrappedRecv conn th istatus (settingsSlowlorisSize settings))
360343 writeIORef istatus True
@@ -429,7 +412,8 @@ serveConnection conn ii th origAddr transport settings app = do
429412 errorResponse e = settingsOnExceptionResponse settings e
430413
431414 http1 firstRequest addr istatus src = do
432− (req, mremainingRef, idxhdr, nextBodyFlush) <- recvRequest firstRequest settings conn ii th addr src transport
415+ (req', mremainingRef, idxhdr, nextBodyFlush) <- recvRequest firstRequest settings conn ii th addr src
416+ let req = req' { isSecure = isTransportSecure transport }
433417 keepAlive <- processRequest istatus src req mremainingRef idxhdr nextBodyFlush
434418 `E.catch` \e -> do
435419 settingsOnException settings (Just req) e