@@ -2,7 +2,7 @@
22{-# LANGUAGE OverloadedStrings #-}
33{-# LANGUAGE ScopedTypeVariables #-}
44
5−module RunSpec (main, spec, withApp, connectTo) where
5+module RunSpec (main, spec, withApp, connectTo, MySocket, msWrite, msRead, msClose) where
66
77import Control.Concurrent (forkIO, killThread, threadDelay)
88import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar)
@@ -11,7 +11,7 @@ import qualified Control.Exception as E
1111import Control.Exception.Lifted (bracket, try, IOException, onException)
1212import Control.Monad (forM_, replicateM_, unless)
1313import Control.Monad.IO.Class (MonadIO, liftIO)
14−import Data.ByteString (ByteString, hPutStr, hGetSome)
14+import Data.ByteString (ByteString)
1515import qualified Data.ByteString as S
1616import Data.ByteString.Builder (byteString)
1717import qualified Data.ByteString.Char8 as S8
@@ -23,7 +23,6 @@ import Network.Socket
2323import Network.Socket.ByteString (sendAll)
2424import Network.Wai hiding (responseHeaders)
2525import Network.Wai.Handler.Warp
26−import System.IO (hFlush, hClose, Handle, IOMode(..))
2726import System.IO.Unsafe (unsafePerformIO)
2827import System.Timeout (timeout)
2928import Test.Hspec
@@ -36,13 +35,44 @@ main = hspec spec
3635type Counter = I.IORef (Either String Int)
3736type CounterApplication = Counter -> Application
3837
39−connectTo :: HostName -> Int -> IO Handle
40−connectTo host port = do
41− let hints = defaultHints { addrSocketType = Stream }
42− addr:_ <- getAddrInfo (Just hints) (Just host) (Just $ show port)
43− sock <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
44− connect sock $ addrAddress addr
45− socketToHandle sock ReadWriteMode
38+data MySocket = MySocket
39+ { msSocket :: !Socket
40+ , _msBuffer :: !(I.IORef ByteString)
41+ }
42+
43+msWrite :: MySocket -> ByteString -> IO ()
44+msWrite ms bs = sendAll (msSocket ms) bs
45+
46+msRead :: MySocket -> Int -> IO ByteString
47+msRead (MySocket s ref) expected = do
48+ bs <- I.readIORef ref
49+ inner (bs:) (S.length bs)
50+ where
51+ inner front total =
52+ case compare total expected of
53+ EQ -> do
54+ I.writeIORef ref mempty
55+ pure $ S.concat $ front []
56+ GT -> do
57+ let bs = S.concat $ front []
58+ (x, y) = S.splitAt expected bs
59+ I.writeIORef ref y
60+ pure x
61+ LT -> do
62+ bs <- safeRecv s 4096
63+ if S.null bs
64+ then do
65+ I.writeIORef ref mempty
66+ pure $ S.concat $ front []
67+ else inner (front . (bs:)) (total + S.length bs)
68+
69+msClose :: MySocket -> IO ()
70+msClose = Network.Socket.close . msSocket
71+
72+connectTo :: Int -> IO MySocket
73+connectTo port = MySocket
74+ <$> (fst <$> getSocketTCP "127.0.0.1" port)
75+ <*> I.newIORef mempty
4676
4777incr :: MonadIO m => Counter -> m ()
4878incr icount = liftIO $ I.atomicModifyIORef icount $ \ecount ->
@@ -115,9 +145,9 @@ runTest :: Int -- ^ expected number of requests
115145runTest expected app chunks = do
116146 ref <- I.newIORef (Right 0)
117147 withApp defaultSettings (app ref) $ \port -> do
118− handle <- connectTo "127.0.0.1" port
119− forM_ chunks $ \chunk -> hPutStr handle chunk >> hFlush handle
120− _ <- timeout 100000 $ replicateM_ expected $ hGetSome handle 4096
148+ ms <- connectTo port
149+ forM_ chunks $ \chunk -> msWrite ms chunk
150+ _ <- timeout 100000 $ replicateM_ expected $ msRead ms 4096
121151 res <- I.readIORef ref
122152 case res of
123153 Left s -> error s
@@ -133,10 +163,9 @@ runTerminateTest expected input = do
133163 ref <- I.newIORef Nothing
134164 let onExc _ = I.writeIORef ref . Just
135165 withApp (setOnException onExc defaultSettings) dummyApp $ \port -> do
136− handle <- connectTo "127.0.0.1" port
137− hPutStr handle input
138− hFlush handle
139− hClose handle
166+ handle <- connectTo port
167+ msWrite handle input
168+ msClose handle
140169 threadDelay 1000
141170 res <- I.readIORef ref
142171 show res `shouldBe` show (Just expected)
@@ -208,13 +237,12 @@ spec = do
208237 liftIO $ I.writeIORef iheaders $ requestHeaders req
209238 f $ responseLBS status200 [] ""
210239 withApp defaultSettings app $ \port -> do
211− handle <- connectTo "127.0.0.1" port
240+ handle <- connectTo port
212241 let input = S.concat
213242 [ "GET / HTTP/1.1\r\nfoo: bar\r\n baz\r\n\tbin\r\n\r\n"
214243 ]
215− hPutStr handle input
216− hFlush handle
217− hClose handle
244+ msWrite handle input
245+ msClose handle
218246 threadDelay 1000
219247 headers <- I.readIORef iheaders
220248 headers `shouldBe`
@@ -226,13 +254,12 @@ spec = do
226254 liftIO $ I.writeIORef iheaders $ requestHeaders req
227255 f $ responseLBS status200 [] ""
228256 withApp defaultSettings app $ \port -> do
229− handle <- connectTo "127.0.0.1" port
257+ handle <- connectTo port
230258 let input = S.concat
231259 [ "GET / HTTP/1.1\r\nfoo:bar\r\n\r\n"
232260 ]
233− hPutStr handle input
234− hFlush handle
235− hClose handle
261+ msWrite handle input
262+ msClose handle
236263 threadDelay 1000
237264 headers <- I.readIORef iheaders
238265 headers `shouldBe`
@@ -249,16 +276,15 @@ spec = do
249276 atomically $ modifyTVar countVar (+ 1)
250277 f $ responseLBS status200 [] ""
251278 withApp defaultSettings app $ \port -> do
252− handle <- connectTo "127.0.0.1" port
279+ handle <- connectTo port
253280 let input = S.concat
254281 [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
255282 , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n\r\n"
256283 , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
257284 , "b\r\nHello World\r\n0\r\n\r\n"
258285 ]
259− hPutStr handle input
260− hFlush handle
261− hClose handle
286+ msWrite handle input
287+ msClose handle
262288 atomically $ do
263289 count <- readTVar countVar
264290 check $ count == 2
@@ -276,13 +302,13 @@ spec = do
276302 atomically $ modifyTVar countVar (+ 1)
277303 f $ responseLBS status200 [] ""
278304 withApp defaultSettings app $ \port -> do
279− handle <- connectTo "127.0.0.1" port
305+ handle <- connectTo port
280306 let input = concat $ replicate 2 $
281307 ["POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"] ++
282308 (replicate 50 "5\r\n12345\r\n") ++
283309 ["0\r\n\r\n"]
284− mapM_ (\bs -> hPutStr handle bs >> hFlush handle) input
285− hClose handle
310+ mapM_ (msWrite handle) input
311+ msClose handle
286312 atomically $ do
287313 count <- readTVar countVar
288314 check $ count == 2
@@ -300,15 +326,15 @@ spec = do
300326 atomically $ modifyTVar countVar (+ 1)
301327 f $ responseLBS status200 [] ""
302328 withApp defaultSettings app $ \port -> do
303− handle <- connectTo "127.0.0.1" port
329+ handle <- connectTo port
304330 let input = S.concat
305331 [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
306332 , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n"
307333 , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
308334 , "b\r\nHello World\r\n0\r\n\r\n"
309335 ]
310− mapM_ (\bs -> hPutStr handle bs >> hFlush handle) $ map S.singleton $ S.unpack input
311− hClose handle
336+ mapM_ (msWrite handle) $ map S.singleton $ S.unpack input
337+ msClose handle
312338 atomically $ do
313339 count <- readTVar countVar
314340 check $ count == 2
@@ -334,16 +360,16 @@ spec = do
334360 let bs1 = S.replicate 2048 88
335361 bs2 = "This is short"
336362 bs = S.append bs1 bs2
337− handle <- connectTo "127.0.0.1" port
338− hPutStr handle "POST / HTTP/1.1\r\n"
339− hPutStr handle "content-length: "
340− hPutStr handle $ S8.pack $ show $ S.length bs
341− hPutStr handle "\r\n\r\n"
363+ handle <- connectTo port
364+ msWrite handle "POST / HTTP/1.1\r\n"
365+ msWrite handle "content-length: "
366+ msWrite handle $ S8.pack $ show $ S.length bs
367+ msWrite handle "\r\n\r\n"
342368 threadDelay 100000
343− hPutStr handle bs1
369+ msWrite handle bs1
344370 threadDelay 100000
345− hPutStr handle bs2
346− hClose handle
371+ msWrite handle bs2
372+ msClose handle
347373 threadDelay 5000000
348374 front <- I.readIORef ifront
349375 S.concat (front []) `shouldBe` bs
@@ -360,13 +386,11 @@ spec = do
360386 loop
361387 doubleBS = S.concatMap $ \w -> S.pack [w, w]
362388 withApp defaultSettings app $ \port -> do
363− handle <- connectTo "127.0.0.1" port
364− hPutStr handle "POST / HTTP/1.1\r\n\r\n12345"
365− hFlush handle
366− timeout 100000 (S.hGet handle 10) >>= (`shouldBe` Just "1122334455")
367− hPutStr handle "67890"
368− hFlush handle
369− timeout 100000 (S.hGet handle 10) >>= (`shouldBe` Just "6677889900")
389+ ms <- connectTo port
390+ msWrite ms "POST / HTTP/1.1\r\n\r\n12345"
391+ timeout 100000 (msRead ms 10) >>= (`shouldBe` Just "1122334455")
392+ msWrite ms "67890"
393+ timeout 100000 (msRead ms 10) >>= (`shouldBe` Just "6677889900")
370394
371395 it "only one date and server header" $ do
372396 let app _ f = f $ responseLBS status200