From fb4afa7aa16641f7642b797d55af3465e6a103fa Mon Sep 17 00:00:00 2001
From: Paolo Capriotti <paolo@capriotti.io>
Date: Fri, 10 Dec 2021 11:33:40 +0100
Subject: [PATCH] Add preface race condition test

This tests that a client always sends the connection preface before any
other frame (see PR #33).
---
 http2.cabal              |  1 +
 test/HTTP2/ServerSpec.hs | 55 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/http2.cabal b/http2.cabal
index ea1adfa5..bffde5ff 100644
--- a/http2.cabal
+++ b/http2.cabal
@@ -170,6 +170,7 @@ Test-Suite spec
                       , hspec >= 1.3
                       , http-types
                       , http2
+                      , network
                       , network-run >= 0.1.0
                       , typed-process
   Default-Extensions:   Strict StrictData
diff --git a/test/HTTP2/ServerSpec.hs b/test/HTTP2/ServerSpec.hs
index 3d4ceec2..a0c3deb7 100644
--- a/test/HTTP2/ServerSpec.hs
+++ b/test/HTTP2/ServerSpec.hs
@@ -13,13 +13,17 @@ import qualified Data.ByteString as B
 import Data.ByteString.Builder (byteString)
 import Data.ByteString.Char8
 import qualified Data.ByteString.Char8 as C8
+import Data.IORef
 import Network.HTTP.Types
 import Network.Run.TCP
+import Network.Socket
+import Network.Socket.ByteString
 import Test.Hspec
 
 import Network.HPACK
 import qualified Network.HTTP2.Client as C
 import Network.HTTP2.Server
+import Network.HTTP2.Frame
 
 port :: String
 port = "8080"
@@ -33,7 +37,14 @@ spec = do
         it "handles normal cases" $
             E.bracket (forkIO runServer) killThread $ \_ -> do
                 threadDelay 10000
-                runClient
+                (runClient allocSimpleConfig)
+        it "should always send the connection preface first" $ do
+            prefaceVar <- newEmptyMVar
+            E.bracket (forkIO (runFakeServer prefaceVar)) killThread $ \_ -> do
+                threadDelay 10000
+                (runClient allocSlowPrefaceConfig)
+            preface <- takeMVar prefaceVar
+            preface `shouldBe` connectionPreface
 
 runServer :: IO ()
 runServer = runTCPServer (Just host) port runHTTP2Server
@@ -42,6 +53,27 @@ runServer = runTCPServer (Just host) port runHTTP2Server
                                  freeSimpleConfig
                                  (`run` server)
 
+runFakeServer :: MVar ByteString -> IO ()
+runFakeServer prefaceVar = do
+  runTCPServer (Just host) port $ \s -> do
+    ref <- newIORef Nothing
+
+    -- send settings
+    sendAll s $ "\x00\x00\x12\x04\x00\x00\x00\x00\x00"
+      `mappend` "\x00\x03\x00\x00\x00\x80\x00\x04\x00"
+      `mappend` "\x01\x00\x00\x00\x05\x00\xff\xff\xff"
+
+    -- receive preface
+    value <- defaultReadN s ref (B.length connectionPreface)
+    putMVar prefaceVar value
+
+    -- send goaway frame
+    sendAll s "\x00\x00\x08\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"
+
+    -- wait for a few ms to make sure the client has a chance to close the
+    -- socket on its end
+    threadDelay 10000
+
 server :: Server
 server req _aux sendResponse = case requestMethod req of
   Just "GET"  -> case requestPath req of
@@ -100,16 +132,31 @@ trailersMaker ctx (Just bs) = return $ NextTrailersMaker $ trailersMaker ctx'
   where
     !ctx' = CH.hashUpdate ctx bs
 
-runClient :: IO ()
-runClient = runTCPClient host port $ runHTTP2Client
+runClient :: (Socket -> BufferSize -> IO Config) -> IO ()
+runClient allocConfig =
+  E.catch (runTCPClient host port $ runHTTP2Client) ignoreHTTP2Error
   where
     authority = C8.pack host
     cliconf = C.ClientConfig "http" authority 20
-    runHTTP2Client s = E.bracket (allocSimpleConfig s 4096)
+    runHTTP2Client s = E.bracket (allocConfig s 4096)
                                  freeSimpleConfig
                                  (\conf -> C.run cliconf conf client)
     client sendRequest = mapConcurrently_ ($ sendRequest) clients
     clients = [client0,client1,client2,client3,client4]
+    ignoreHTTP2Error :: HTTP2Error -> IO ()
+    ignoreHTTP2Error _ = pure ()
+
+-- delay sending preface to be able to test if it is always sent first
+allocSlowPrefaceConfig :: Socket -> BufferSize -> IO Config
+allocSlowPrefaceConfig s size = do
+  config <- allocSimpleConfig s size
+  pure config { confSendAll = slowPrefaceSend (confSendAll config) }
+  where
+    slowPrefaceSend :: (ByteString -> IO ()) -> ByteString -> IO ()
+    slowPrefaceSend orig chunk = do
+      when (C8.pack "PRI" `isPrefixOf` chunk) $ do
+        threadDelay 10000
+      orig chunk
 
 client0 :: C.Client ()
 client0 sendRequest = do