git.iko.soy/typeable/wai
on master help sign in
Merge pull request #392 from yesodweb/input-buffer

Input buffer

commitbffaa1a296f28abb8c18fcacc22c069de7ce12fc authorKazu Yamamoto <kazu@iij.ad.jp> dateMon 29 Jun 2015, 20:27 (11 years ago) parent69c30dc281340d
5 files changed+113−49
compare with master
unifiedsplit
warp-tls/Network/Wai/Handler/WarpTLS.hs+13−31 warp/Network/Wai/Handler/Warp/Buffer.hs+13−11 warp/Network/Wai/Handler/Warp/Internal.hs+6 warp/Network/Wai/Handler/Warp/Recv.hs+72−5 warp/Network/Wai/Handler/Warp/Types.hs+9−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+13−11
@@ -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 ()
@@ -102,11 +102,13 @@ toBlazeBuffer ptr size = do
102102 fptr <- newForeignPtr_ ptr
103103 return $ B.Buffer fptr ptr ptr (ptr `plusPtr` size)
104104
105−{-# INLINE copy #-}
106−copy :: Ptr Word8 -> ByteString -> IO (Ptr Word8)
105+-- | Copying the bytestring to the buffer.
106+-- This function returns the point where the next copy should start.
107+copy :: Buffer -> ByteString -> IO Buffer
107108copy !ptr (PS fp o l) = withForeignPtr fp $ \p -> do
108109 memcpy ptr (p `plusPtr` o) (fromIntegral l)
109110 return $! ptr `plusPtr` l
111+{-# INLINE copy #-}
110112
111113{-# INLINE toBS #-}
112114toBS :: Buffer -> Int -> IO ByteString
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+72−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,72 @@ 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+-- | This function returns a receiving function
42+-- based on two receiving functions.
43+-- The returned function efficiently manages received data
44+-- which is initialized by the first argument.
45+-- The returned function may allocate a byte string with malloc().
46+makePlainReceiveN :: Socket -> ByteString -> IO (BufSize -> IO ByteString)
47+makePlainReceiveN s bs0 = do
48+ ref <- newIORef bs0
49+ pool <- newBufferPool
50+ return $ receiveN ref (receive s pool) (receiveBuf s)
51+
52+receiveN :: IORef ByteString -> Recv -> RecvBuf -> BufSize -> IO ByteString
53+receiveN ref recv recvBuf size = do
54+ cached <- readIORef ref
55+ (bs, leftover) <- spell cached size recv recvBuf
56+ writeIORef ref leftover
57+ return bs
58+
59+----------------------------------------------------------------
60+
61+spell :: ByteString -> BufSize -> IO ByteString -> RecvBuf -> IO (ByteString, ByteString)
62+spell initial siz recv recvBuf
63+ | siz <= len0 = return $ BS.splitAt siz initial
64+ | siz <= 4096 = do
65+ bs <- recv
66+ if bs == "" then
67+ return ("", "")
68+ else do
69+ let (bs1, leftover) = BS.splitAt (siz - len0) bs
70+ return (BS.append initial bs1, leftover)
71+ | otherwise = do
72+ bs@(PS fptr _ _) <- mallocByteString siz
73+ withForeignPtr fptr $ \ptr -> do
74+ ptr' <- copy ptr initial
75+ full <- recvBuf ptr' (siz - len0)
76+ if full then
77+ return (bs, "")
78+ else
79+ return ("", "") -- fixme
80+ where
81+ len0 = BS.length initial
82+
83+receive :: Socket -> BufferPool -> Recv
3084receive sock pool = withBufferPool pool $ \ (ptr, size) -> do
3185 let sock' = fdSocket sock
3286 size' = fromIntegral size
3387 fromIntegral <$> receiveloop sock' ptr size'
3488
89+receiveBuf :: Socket -> RecvBuf
90+receiveBuf sock buf0 siz0 = loop buf0 siz0
91+ where
92+ loop _ 0 = return True
93+ loop buf siz = do
94+ n <- fromIntegral <$> receiveloop fd buf (fromIntegral siz)
95+ -- fixme: what should we do in the case of n == 0
96+ if n == 0 then
97+ return False
98+ else
99+ loop (buf `plusPtr` n) (siz - n)
100+ fd = fdSocket sock
101+
35102receiveloop :: CInt -> Ptr Word8 -> CSize -> IO CInt
36103receiveloop sock ptr size = do
37104#ifdef mingw32_HOST_OS
warp/Network/Wai/Handler/Warp/Types.hs+9−2
@@ -95,6 +95,13 @@ type Buffer = Ptr Word8
9595-- | Type for buffer size
9696type BufSize = Int
9797
98+-- | Type for the action to receive input data
99+type Recv = IO ByteString
100+
101+-- | Type for the action to receive input data with a buffer.
102+-- The result boolean indicates whether or not the buffer is fully filled.
103+type RecvBuf = Buffer -> BufSize -> IO Bool
104+
98105-- | Data type to manipulate IO actions for connections.
99106-- This is used to abstract IO actions for plain HTTP and HTTP over TLS.
100107data Connection = Connection {
@@ -106,8 +113,8 @@ data Connection = Connection {
106113 , connSendFile :: SendFile
107114 -- | The connection closing function.
108115 , connClose :: IO ()
109− -- | The connection receiving function.
110− , connRecv :: IO ByteString
116+ -- | The connection receiving function. This returns "" for EOF.
117+ , connRecv :: Recv
111118 -- | The write buffer.
112119 , connWriteBuffer :: Buffer
113120 -- | The size of the write buffer.