content: {-# LANGUAGE OverloadedStrings #-}
  
  module HandshakeSpec where
  
  import Control.Concurrent
  import qualified Control.Exception as E
  import Control.Monad
  import qualified Data.ByteString as BS
  import Network.TLS (Group (..), HandshakeMode13 (..))
  import qualified Network.TLS as TLS
  import Test.Hspec
  
  import Network.QUIC
  import Network.QUIC.Client as C
  import Network.QUIC.Internal hiding (RTT0)
  import Network.QUIC.Server as S
  
  import Config
  
  spec :: Spec
  spec = do
      sc0' <- runIO makeTestServerConfig
      smgr <- runIO newSessionManager
      var <- runIO newEmptyMVar
      let sc0 =
              sc0'
                  { scSessionManager = smgr
                  , scHooks =
                      (scHooks sc0')
                          { onServerReady = putMVar var ()
                          }
                  }
      let waitS = takeMVar var :: IO ()
      describe "handshake" $ do
          it "can handshake in the normal case" $ do
              let cc = testClientConfig
                  sc = sc0
              testHandshake cc sc waitS FullHandshake
          it "can handshake in the case of TLS hello retry" $ do
              let cc = testClientConfig
                  sc = sc0{scGroups = [P256]}
              testHandshake cc sc waitS HelloRetryRequest
          it "can handshake in the case of QUIC retry" $ do
              let cc = testClientConfig
                  sc = sc0{scRequireRetry = True}
              testHandshake cc sc waitS FullHandshake
          it "can handshake in the case of resumption" $ do
              let cc = testClientConfig
                  sc = sc0
              testHandshake2 cc sc waitS (FullHandshake, PreSharedKey) False
          it "can handshake in the case of 0-RTT" $ do
              let cc = testClientConfig
                  sc = sc0{scUse0RTT = True}
              testHandshake2 cc sc waitS (FullHandshake, RTT0) True
          it "fails with unknown server certificate" $ do
              let cc1 =
                      testClientConfig
                          { ccValidate = True -- ouch, default should be reversed
                          }
                  cc2 = testClientConfig
                  sc = sc0
                  certificateRejected e
                      | TransportErrorIsSent te@(TransportError _) _ <- e =
                          te == cryptoError TLS.CertificateUnknown
                      | otherwise = False
              testHandshake3 cc1 cc2 sc waitS certificateRejected
          it "fails with no group in common" $ do
              let cc1 = testClientConfig{ccGroups = [X25519]}
                  cc2 = testClientConfig{ccGroups = [P256]}
                  sc = sc0{scGroups = [P256]}
                  handshakeFailure e
                      | TransportErrorIsReceived te@(TransportError _) _ <- e =
                          te == cryptoError TLS.HandshakeFailure
                      | otherwise = False
              testHandshake3 cc1 cc2 sc waitS handshakeFailure
          it "can handshake with large HE from a client" $ do
              let cc0 = testClientConfig
                  params =
                      (ccParameters cc0)
                          { grease = Just (BS.pack (replicate 2400 0))
                          }
                  cc = cc0{ccParameters = params}
                  sc = sc0
              testHandshake cc sc waitS FullHandshake
          it "can handshake with large EE from a server (3-times rule)" $ do
              let cc = testClientConfig
                  params =
                      (scParameters sc0)
                          { grease = Just (BS.pack (replicate 3800 0))
                          }
                  sc = sc0{scParameters = params}
              testHandshake cc sc waitS FullHandshake
  
  onE :: IO b -> IO a -> IO a
  onE h b = E.onException b h
  
  testHandshake
      :: ClientConfig -> ServerConfig -> IO () -> HandshakeMode13 -> IO ()
  testHandshake cc sc waitS mode = do
      mvar <- newEmptyMVar
      E.bracket (forkIO $ server mvar) killThread $ \_ -> client mvar
    where
      client mvar = do
          waitS
          C.run cc $ \conn -> do
              waitEstablished conn
              handshakeMode <$> getConnectionInfo conn `shouldReturn` mode
              takeMVar mvar
      server mvar = S.run sc $ \conn -> do
          waitEstablished conn
          handshakeMode <$> getConnectionInfo conn `shouldReturn` mode
          putMVar mvar ()
  
  query :: BS.ByteString -> Connection -> IO ()
  query content conn = do
      waitEstablished conn
      s <- stream conn
      sendStream s content
      shutdownStream s
      void $ recvStream s 1024
  
  testHandshake2
      :: ClientConfig
      -> ServerConfig
      -> IO ()
      -> (HandshakeMode13, HandshakeMode13)
      -> Bool
      -> IO ()
  testHandshake2 cc1 sc waitS (mode1, mode2) use0RTT = do
      mvar <- newEmptyMVar
      E.bracket (forkIO $ server mvar) killThread $ \_ -> client mvar
    where
      runClient cc mode action = C.run cc $ \conn -> do
          void $ action conn
          handshakeMode <$> getConnectionInfo conn `shouldReturn` mode
          threadDelay 50000
          getResumptionInfo conn
      client mvar = do
          waitS
          res <- runClient cc1 mode1 $ query "first"
          threadDelay 50000
          let cc2 =
                  cc1
                      { ccResumption = res
                      , ccUse0RTT = use0RTT
                      }
          void $ runClient cc2 mode2 $ query "second"
          takeMVar mvar
      server mvar = S.run sc serv
        where
          serv conn = do
              s <- acceptStream conn
              bs <- recvStream s 1024
              sendStream s "bye"
              closeStream s
              when (bs == "second") $ putMVar mvar ()
  
  testHandshake3
      :: ClientConfig
      -> ClientConfig
      -> ServerConfig
      -> IO ()
      -> (QUICException -> Bool)
      -> IO ()
  testHandshake3 cc1 cc2 sc waitS selector = do
      mvar <- newEmptyMVar
      E.bracket (forkIO $ server mvar) killThread $ \_ -> client mvar
    where
      client mvar = do
          waitS
          C.run cc1 (query "first") `shouldThrow` selector
          C.run cc2 (query "second") `shouldReturn` ()
          takeMVar mvar
      server mvar = S.run sc $ \conn -> do
          s <- acceptStream conn
          recvStream s 1024 `shouldReturn` "second"
          sendStream s "bye"
          closeStream s
          putMVar mvar ()
  