@@ -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 ->
@@ -106,7 +136,13 @@ withApp settings app f = do
106136 bracket
107137 (forkIO $ runSettings settings' app `onException` putMVar baton ())
108138 killThread
109− (const $ takeMVar baton >> f port)
139+ (const $ do
140+ takeMVar baton
141+ -- use timeout to make sure we don't take too long
142+ mres <- timeout (60 * 1000 * 1000) (f port)
143+ case mres of
144+ Nothing -> error "Timeout triggered, too slow!"
145+ Just a -> pure a)
110146
111147runTest :: Int -- ^ expected number of requests
112148 -> CounterApplication
@@ -115,9 +151,9 @@ runTest :: Int -- ^ expected number of requests
115151runTest expected app chunks = do
116152 ref <- I.newIORef (Right 0)
117153 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
154+ ms <- connectTo port
155+ forM_ chunks $ \chunk -> msWrite ms chunk
156+ _ <- timeout 100000 $ replicateM_ expected $ msRead ms 4096
121157 res <- I.readIORef ref
122158 case res of
123159 Left s -> error s
@@ -133,10 +169,9 @@ runTerminateTest expected input = do
133169 ref <- I.newIORef Nothing
134170 let onExc _ = I.writeIORef ref . Just
135171 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
172+ handle <- connectTo port
173+ msWrite handle input
174+ msClose handle
140175 threadDelay 1000
141176 res <- I.readIORef ref
142177 show res `shouldBe` show (Just expected)
@@ -208,13 +243,12 @@ spec = do
208243 liftIO $ I.writeIORef iheaders $ requestHeaders req
209244 f $ responseLBS status200 [] ""
210245 withApp defaultSettings app $ \port -> do
211− handle <- connectTo "127.0.0.1" port
246+ handle <- connectTo port
212247 let input = S.concat
213248 [ "GET / HTTP/1.1\r\nfoo: bar\r\n baz\r\n\tbin\r\n\r\n"
214249 ]
215− hPutStr handle input
216− hFlush handle
217− hClose handle
250+ msWrite handle input
251+ msClose handle
218252 threadDelay 1000
219253 headers <- I.readIORef iheaders
220254 headers `shouldBe`
@@ -226,13 +260,12 @@ spec = do
226260 liftIO $ I.writeIORef iheaders $ requestHeaders req
227261 f $ responseLBS status200 [] ""
228262 withApp defaultSettings app $ \port -> do
229− handle <- connectTo "127.0.0.1" port
263+ handle <- connectTo port
230264 let input = S.concat
231265 [ "GET / HTTP/1.1\r\nfoo:bar\r\n\r\n"
232266 ]
233− hPutStr handle input
234− hFlush handle
235− hClose handle
267+ msWrite handle input
268+ msClose handle
236269 threadDelay 1000
237270 headers <- I.readIORef iheaders
238271 headers `shouldBe`
@@ -241,28 +274,33 @@ spec = do
241274
242275 describe "chunked bodies" $ do
243276 it "works" $ do
277+ countVar <- newTVarIO (0 :: Int)
244278 ifront <- I.newIORef id
245279 let app req f = do
246280 bss <- consumeBody $ requestBody req
247281 liftIO $ I.atomicModifyIORef ifront $ \front -> (front . (S.concat bss:), ())
282+ atomically $ modifyTVar countVar (+ 1)
248283 f $ responseLBS status200 [] ""
249284 withApp defaultSettings app $ \port -> do
250− handle <- connectTo "127.0.0.1" port
285+ handle <- connectTo port
251286 let input = S.concat
252287 [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
253288 , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n\r\n"
254289 , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
255290 , "b\r\nHello World\r\n0\r\n\r\n"
256291 ]
257− hPutStr handle input
258− hFlush handle
259− hClose handle
260− threadDelay 1000
292+ msWrite handle input
293+ msClose handle
294+ atomically $ do
295+ count <- readTVar countVar
296+ check $ count == 2
261297 front <- I.readIORef ifront
262298 front [] `shouldBe`
263299 [ "Hello World\nBye"
264300 , "Hello World"
265301 ]
302+#if !WINDOWS
303+-- Too slow on Windows
266304 it "lots of chunks" $ do
267305 ifront <- I.newIORef id
268306 countVar <- newTVarIO (0 :: Int)
@@ -272,13 +310,13 @@ spec = do
272310 atomically $ modifyTVar countVar (+ 1)
273311 f $ responseLBS status200 [] ""
274312 withApp defaultSettings app $ \port -> do
275− handle <- connectTo "127.0.0.1" port
313+ handle <- connectTo port
276314 let input = concat $ replicate 2 $
277315 ["POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"] ++
278316 (replicate 50 "5\r\n12345\r\n") ++
279317 ["0\r\n\r\n"]
280− mapM_ (\bs -> hPutStr handle bs >> hFlush handle) input
281− hClose handle
318+ mapM_ (msWrite handle) input
319+ msClose handle
282320 atomically $ do
283321 count <- readTVar countVar
284322 check $ count == 2
@@ -286,7 +324,6 @@ spec = do
286324 front [] `shouldBe` replicate 2 (S.concat $ replicate 50 "12345")
287325-- For some reason, the following test on Windows causes the socket
288326-- to be killed prematurely. Worth investigating in the future if possible.
289−#if !WINDOWS
290327 it "in chunks" $ do
291328 ifront <- I.newIORef id
292329 countVar <- newTVarIO (0 :: Int)
@@ -296,15 +333,15 @@ spec = do
296333 atomically $ modifyTVar countVar (+ 1)
297334 f $ responseLBS status200 [] ""
298335 withApp defaultSettings app $ \port -> do
299− handle <- connectTo "127.0.0.1" port
336+ handle <- connectTo port
300337 let input = S.concat
301338 [ "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
302339 , "c\r\nHello World\n\r\n3\r\nBye\r\n0\r\n"
303340 , "POST / HTTP/1.1\r\nTransfer-Encoding: Chunked\r\n\r\n"
304341 , "b\r\nHello World\r\n0\r\n\r\n"
305342 ]
306− mapM_ (\bs -> hPutStr handle bs >> hFlush handle) $ map S.singleton $ S.unpack input
307− hClose handle
343+ mapM_ (msWrite handle) $ map S.singleton $ S.unpack input
344+ msClose handle
308345 atomically $ do
309346 count <- readTVar countVar
310347 check $ count == 2
@@ -330,16 +367,16 @@ spec = do
330367 let bs1 = S.replicate 2048 88
331368 bs2 = "This is short"
332369 bs = S.append bs1 bs2
333− handle <- connectTo "127.0.0.1" port
334− hPutStr handle "POST / HTTP/1.1\r\n"
335− hPutStr handle "content-length: "
336− hPutStr handle $ S8.pack $ show $ S.length bs
337− hPutStr handle "\r\n\r\n"
370+ handle <- connectTo port
371+ msWrite handle "POST / HTTP/1.1\r\n"
372+ msWrite handle "content-length: "
373+ msWrite handle $ S8.pack $ show $ S.length bs
374+ msWrite handle "\r\n\r\n"
338375 threadDelay 100000
339− hPutStr handle bs1
376+ msWrite handle bs1
340377 threadDelay 100000
341− hPutStr handle bs2
342− hClose handle
378+ msWrite handle bs2
379+ msClose handle
343380 threadDelay 5000000
344381 front <- I.readIORef ifront
345382 S.concat (front []) `shouldBe` bs
@@ -356,30 +393,33 @@ spec = do
356393 loop
357394 doubleBS = S.concatMap $ \w -> S.pack [w, w]
358395 withApp defaultSettings app $ \port -> do
359− handle <- connectTo "127.0.0.1" port
360− hPutStr handle "POST / HTTP/1.1\r\n\r\n12345"
361− hFlush handle
362− timeout 100000 (S.hGet handle 10) >>= (`shouldBe` Just "1122334455")
363− hPutStr handle "67890"
364− hFlush handle
365− timeout 100000 (S.hGet handle 10) >>= (`shouldBe` Just "6677889900")
396+ ms <- connectTo port
397+ msWrite ms "POST / HTTP/1.1\r\n\r\n12345"
398+ timeout 100000 (msRead ms 10) >>= (`shouldBe` Just "1122334455")
399+ msWrite ms "67890"
400+ timeout 100000 (msRead ms 10) >>= (`shouldBe` Just "6677889900")
366401
367402 it "only one date and server header" $ do
368403 let app _ f = f $ responseLBS status200
369404 [ ("server", "server")
370405 , ("date", "date")
371406 ] ""
407+ getValues key = map snd
408+ . filter (\(key', _) -> key == key')
409+ . responseHeaders
372410 withApp defaultSettings app $ \port -> do
373411 res <- sendGET $ "http://127.0.0.1:" ++ show port
374− getHeaderValue hServer (responseHeaders res) `shouldBe` Just "server"
375− getHeaderValue hDate (responseHeaders res) `shouldBe` Just "date"
412+ getValues hServer res `shouldBe` ["server"]
413+ getValues hDate res `shouldBe` ["date"]
376414
377415 it "streaming echo #249" $ do
416+ countVar <- newTVarIO (0 :: Int)
378417 let app req f = f $ responseStream status200 [] $ \write _ -> do
379418 let loop = do
380419 bs <- requestBody req
381420 unless (S.null bs) $ do
382421 write $ byteString bs
422+ atomically $ modifyTVar countVar (+ 1)
383423 loop
384424 loop
385425 withApp defaultSettings app $ \port -> do
@@ -387,6 +427,9 @@ spec = do
387427 sendAll sock "POST / HTTP/1.1\r\ntransfer-encoding: chunked\r\n\r\n"
388428 threadDelay 10000
389429 sendAll sock "5\r\nhello\r\n0\r\n\r\n"
430+ atomically $ do
431+ count <- readTVar countVar
432+ check $ count >= 1
390433 bs <- safeRecv sock 4096
391434 S.takeWhile (/= 13) bs `shouldBe` "HTTP/1.1 200 OK"
392435