git.iko.soy/typeable/wai
on master help sign in
improve buffer management of plain input in WarpTLS.
commitcb4cd35873cca49bd3f123fe39b9a87973c0019d authorKazu Yamamoto <kazu@iij.ad.jp> dateMon 29 Jun 2015, 14:08 (11 years ago) parent69c30dc
5 files changed+101−48
compare with master
unifiedsplit
warp-tls/Network/Wai/Handler/WarpTLS.hs+13−31 warp/Network/Wai/Handler/Warp/Buffer.hs+10−10 warp/Network/Wai/Handler/Warp/Internal.hs+6 warp/Network/Wai/Handler/Warp/Recv.hs+67−5 warp/Network/Wai/Handler/Warp/Types.hs+5−2
warp-tls/Network/Wai/Handler/WarpTLS.hs+13−31
@@ -50,6 +50,7 @@ import qualified Network.TLS as TLS
5050import qualified Network.TLS.Extra as TLSExtra
5151import Network.Wai (Application)
5252import Network.Wai.Handler.Warp
53+import Network.Wai.Handler.Warp.Internal
5354import System.IO.Error (isEOFError)
5455
5556----------------------------------------------------------------
@@ -255,21 +256,21 @@ getter tlsset@TLSSettings{..} sock params = do
255256mkConn :: TLS.TLSParams params => TLSSettings -> Socket -> params -> IO (Connection, Transport)
256257mkConn tlsset s params = do
257258 firstBS <- safeRecv s 4096
258− cachedRef <- I.newIORef firstBS
259259 (if not (S.null firstBS) && S.head firstBS == 0x16 then
260− httpOverTls tlsset s cachedRef params
260+ httpOverTls tlsset s firstBS params
261261 else
262− plainHTTP tlsset s cachedRef) `onException` sClose s
262+ plainHTTP tlsset s firstBS) `onException` sClose s
263263
264264----------------------------------------------------------------
265265
266−httpOverTls :: TLS.TLSParams params => TLSSettings -> Socket -> I.IORef S.ByteString -> params -> IO (Connection, Transport)
267−httpOverTls TLSSettings{..} s cachedRef params = do
266+httpOverTls :: TLS.TLSParams params => TLSSettings -> Socket -> S.ByteString -> params -> IO (Connection, Transport)
267+httpOverTls TLSSettings{..} s bs0 params = do
268+ recvN <- makePlainReceiveN s bs0
268269#if MIN_VERSION_tls(1,3,0)
269− ctx <- TLS.contextNew backend params
270+ ctx <- TLS.contextNew (backend recvN) params
270271#else
271272 gen <- Crypto.Random.AESCtr.makeSystem
272− ctx <- TLS.contextNew backend params gen
273+ ctx <- TLS.contextNew (backend recvN) params gen
273274#endif
274275 TLS.contextHookSetLogging ctx tlsLogging
275276 TLS.handshake ctx
@@ -277,11 +278,11 @@ httpOverTls TLSSettings{..} s cachedRef params = do
277278 tls <- getTLSinfo ctx
278279 return (conn ctx writeBuf, tls)
279280 where
280− backend = TLS.Backend {
281+ backend recvN = TLS.Backend {
281282 TLS.backendFlush = return ()
282283 , TLS.backendClose = sClose s
283284 , TLS.backendSend = sendAll s
284− , TLS.backendRecv = recvTLS cachedRef s
285+ , TLS.backendRecv = recvN
285286 }
286287 conn ctx writeBuf = Connection {
287288 connSendMany = TLS.sendData ctx . L.fromChunks
@@ -338,10 +339,11 @@ tryIO = try
338339
339340----------------------------------------------------------------
340341
341−plainHTTP :: TLSSettings -> Socket -> I.IORef S.ByteString -> IO (Connection, Transport)
342−plainHTTP TLSSettings{..} s cachedRef = case onInsecure of
342+plainHTTP :: TLSSettings -> Socket -> S.ByteString -> IO (Connection, Transport)
343+plainHTTP TLSSettings{..} s bs0 = case onInsecure of
343344 AllowInsecure -> do
344345 conn' <- socketConnection s
346+ cachedRef <- I.newIORef bs0
345347 let conn'' = conn'
346348 { connRecv = recvPlain cachedRef (connRecv conn')
347349 }
@@ -354,26 +356,6 @@ plainHTTP TLSSettings{..} s cachedRef = case onInsecure of
354356
355357----------------------------------------------------------------
356358
357−recvTLS :: I.IORef S.ByteString -> Socket -> Int -> IO S.ByteString
358−recvTLS cachedRef s size = do
359− cached <- I.readIORef cachedRef
360− loop cached
361− where
362− loop bs | S.length bs >= size = do
363− let (x, y) = S.splitAt size bs
364− I.writeIORef cachedRef y
365− return x
366− loop bs1 = do
367− bs2 <- safeRecv s 4096
368− if S.null bs2 then do
369− -- FIXME does this deserve an exception being thrown?
370− I.writeIORef cachedRef S.empty
371− return bs1
372− else
373− loop $ S.append bs1 bs2
374−
375−----------------------------------------------------------------
376−
377359-- | Modify the given receive function to first check the given @IORef@ for a
378360-- chunk of data. If present, takes the chunk of data from the @IORef@ and
379361-- empties out the @IORef@. Otherwise, calls the supplied receive function.
warp/Network/Wai/Handler/Warp/Buffer.hs+10−10
@@ -1,9 +1,10 @@
1−{-# LANGUAGE BangPatterns #-}
1+{-# LANGUAGE BangPatterns, OverloadedStrings #-}
22
33module Network.Wai.Handler.Warp.Buffer (
44 bufferSize
55 , allocateBuffer
66 , freeBuffer
7+ , mallocByteString
78 , newBufferPool
89 , withBufferPool
910 , toBlazeBuffer
@@ -13,14 +14,13 @@ module Network.Wai.Handler.Warp.Buffer (
1314
1415import Control.Monad (when)
1516import qualified Data.ByteString as BS
16−import Data.ByteString.Internal
17+import Data.ByteString.Internal (ByteString(..), memcpy)
1718import Data.ByteString.Unsafe (unsafeTake, unsafeDrop)
1819import Data.IORef (newIORef, readIORef, writeIORef)
1920import qualified Data.Streaming.ByteString.Builder.Buffer as B (Buffer (..))
20−import Data.Word (Word8)
2121import Foreign.ForeignPtr
2222import Foreign.Marshal.Alloc (mallocBytes, free, finalizerFree)
23−import Foreign.Ptr (Ptr, castPtr, plusPtr)
23+import Foreign.Ptr (castPtr, plusPtr)
2424import Network.Wai.Handler.Warp.Types
2525
2626----------------------------------------------------------------
@@ -51,12 +51,12 @@ minBufferSize = 2048
5151newBufferPool :: IO BufferPool
5252newBufferPool = newIORef BS.empty
5353
54−mallocBuffer :: Int -> IO ByteString
55−mallocBuffer size = do
56− ptr <- mallocBytes size
54+mallocByteString :: Int -> IO ByteString
55+mallocByteString size = do
56+ ptr <- allocateBuffer size
5757 fptr <- newForeignPtr finalizerFree ptr
5858 return $! PS fptr 0 size
59−{-# INLINE mallocBuffer #-}
59+{-# INLINE mallocByteString #-}
6060
6161{-
6262createBuffer :: Int -> IO ByteString
@@ -73,7 +73,7 @@ usefulBuffer buffer = BS.length buffer >= minBufferSize
7373getBuffer :: BufferPool -> IO ByteString
7474getBuffer pool = do
7575 buffer <- readIORef pool
76− if usefulBuffer buffer then return buffer else mallocBuffer largeBufferSize
76+ if usefulBuffer buffer then return buffer else mallocByteString largeBufferSize
7777{-# INLINE getBuffer #-}
7878
7979putBuffer :: BufferPool -> ByteString -> IO ()
@@ -103,7 +103,7 @@ toBlazeBuffer ptr size = do
103103 return $ B.Buffer fptr ptr ptr (ptr `plusPtr` size)
104104
105105{-# INLINE copy #-}
106−copy :: Ptr Word8 -> ByteString -> IO (Ptr Word8)
106+copy :: Buffer -> ByteString -> IO Buffer
107107copy !ptr (PS fp o l) = withForeignPtr fp $ \p -> do
108108 memcpy ptr (p `plusPtr` o) (fromIntegral l)
109109 return $! ptr `plusPtr` l
warp/Network/Wai/Handler/Warp/Internal.hs+6
@@ -12,12 +12,17 @@ module Network.Wai.Handler.Warp.Internal (
1212 -- * Connection
1313 , Connection (..)
1414 , socketConnection
15+ -- ** Receive
16+ , Recv
17+ , RecvBuf
18+ , makePlainReceiveN
1519 -- ** Buffer
1620 , Buffer
1721 , BufSize
1822 , bufferSize
1923 , allocateBuffer
2024 , freeBuffer
25+ , copy
2126 -- ** Sendfile
2227 , FileId (..)
2328 , SendFile
@@ -63,6 +68,7 @@ import Network.Wai.Handler.Warp.Buffer
6368import Network.Wai.Handler.Warp.Date
6469import Network.Wai.Handler.Warp.FdCache
6570import Network.Wai.Handler.Warp.Header
71+import Network.Wai.Handler.Warp.Recv
6672import Network.Wai.Handler.Warp.Request
6773import Network.Wai.Handler.Warp.Response
6874import Network.Wai.Handler.Warp.Run
warp/Network/Wai/Handler/Warp/Recv.hs+67−5
@@ -1,22 +1,29 @@
1−{-# LANGUAGE ForeignFunctionInterface #-}
1+{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings #-}
22{-# LANGUAGE CPP #-}
33
44module Network.Wai.Handler.Warp.Recv (
55 receive
6+ , receiveBuf
7+ , makeReceiveN
8+ , makePlainReceiveN
9+ , spell
610 ) where
711
812#if __GLASGOW_HASKELL__ < 709
913import Control.Applicative ((<$>))
1014#endif
11−import Data.ByteString (ByteString)
15+import qualified Data.ByteString as BS
16+import Data.ByteString.Internal (ByteString(..))
1217import Data.Word (Word8)
1318import Foreign.C.Error (eAGAIN, getErrno, throwErrno)
1419import Foreign.C.Types
15−import Foreign.Ptr (Ptr, castPtr)
20+import Foreign.ForeignPtr (withForeignPtr)
21+import Foreign.Ptr (Ptr, castPtr, plusPtr)
1622import GHC.Conc (threadWaitRead)
1723import Network.Socket (Socket, fdSocket)
18−import Network.Wai.Handler.Warp.Types
1924import Network.Wai.Handler.Warp.Buffer
25+import Network.Wai.Handler.Warp.IORef
26+import Network.Wai.Handler.Warp.Types
2027import System.Posix.Types (Fd(..))
2128
2229#ifdef mingw32_HOST_OS
@@ -26,12 +33,67 @@ import Network.Wai.Handler.Warp.Windows
2633
2734----------------------------------------------------------------
2835
29−receive :: Socket -> BufferPool -> IO ByteString
36+makeReceiveN :: ByteString -> Recv -> RecvBuf -> IO (BufSize -> IO ByteString)
37+makeReceiveN bs0 recv recvBuf = do
38+ ref <- newIORef bs0
39+ return $ receiveN ref recv recvBuf
40+
41+makePlainReceiveN :: Socket -> ByteString -> IO (BufSize -> IO ByteString)
42+makePlainReceiveN s bs0 = do
43+ ref <- newIORef bs0
44+ pool <- newBufferPool
45+ return $ receiveN ref (receive s pool) (receiveBuf s)
46+
47+receiveN :: IORef ByteString -> Recv -> RecvBuf -> BufSize -> IO ByteString
48+receiveN ref recv recvBuf size = do
49+ cached <- readIORef ref
50+ (bs, leftover) <- spell cached size recv recvBuf
51+ writeIORef ref leftover
52+ return bs
53+
54+----------------------------------------------------------------
55+
56+spell :: ByteString -> BufSize -> IO ByteString -> RecvBuf -> IO (ByteString, ByteString)
57+spell initial siz recv recvBuf
58+ | siz <= len0 = return $ BS.splitAt siz initial
59+ | siz <= 4096 = do
60+ bs <- recv
61+ if bs == "" then
62+ return ("", "")
63+ else do
64+ let (bs1, leftover) = BS.splitAt (siz - len0) bs
65+ return (BS.append initial bs1, leftover)
66+ | otherwise = do
67+ bs@(PS fptr _ _) <- mallocByteString siz
68+ withForeignPtr fptr $ \ptr -> do
69+ ptr' <- copy ptr initial
70+ full <- recvBuf ptr' (siz - len0)
71+ if full then
72+ return (bs, "")
73+ else
74+ return ("", "") -- fixme
75+ where
76+ len0 = BS.length initial
77+
78+receive :: Socket -> BufferPool -> Recv
3079receive sock pool = withBufferPool pool $ \ (ptr, size) -> do
3180 let sock' = fdSocket sock
3281 size' = fromIntegral size
3382 fromIntegral <$> receiveloop sock' ptr size'
3483
84+receiveBuf :: Socket -> RecvBuf
85+receiveBuf sock buf0 siz0 = loop buf0 siz0
86+ where
87+ loop _ 0 = return True
88+ loop buf siz = do
89+ n <- fromIntegral <$> receiveloop fd buf (fromIntegral siz)
90+ -- fixme: what should we do in the case of n == 0
91+ if n == 0 then
92+ return False
93+ else
94+ loop (buf `plusPtr` n) (siz - n)
95+ fd = fdSocket sock
96+
3597receiveloop :: CInt -> Ptr Word8 -> CSize -> IO CInt
3698receiveloop sock ptr size = do
3799#ifdef mingw32_HOST_OS
warp/Network/Wai/Handler/Warp/Types.hs+5−2
@@ -95,6 +95,9 @@ type Buffer = Ptr Word8
9595-- | Type for buffer size
9696type BufSize = Int
9797
98+type Recv = IO ByteString
99+type RecvBuf = Buffer -> BufSize -> IO Bool
100+
98101-- | Data type to manipulate IO actions for connections.
99102-- This is used to abstract IO actions for plain HTTP and HTTP over TLS.
100103data Connection = Connection {
@@ -106,8 +109,8 @@ data Connection = Connection {
106109 , connSendFile :: SendFile
107110 -- | The connection closing function.
108111 , connClose :: IO ()
109− -- | The connection receiving function.
110− , connRecv :: IO ByteString
112+ -- | The connection receiving function. This returns "" for EOF.
113+ , connRecv :: Recv
111114 -- | The write buffer.
112115 , connWriteBuffer :: Buffer
113116 -- | The size of the write buffer.