xftp: close connection after failed handshake

src/Simplex/FileTransfer/Client.hs

@@ -149,7 +149,8 @@ getXFTPClient transportSession@(_, srv, _) config@XFTPClientConfig {clientALPN,
   thParams@THandleParams {thVersion} <- case sessionALPN of
     Just alpn
       | alpn == xftpALPNv1 || alpn == httpALPN11 ->
-          xftpClientHandshakeV1 serverVRange keyHash http2Client mkEntitlementProof thParams0
+          ExceptT (runExceptT (xftpClientHandshakeV1 serverVRange keyHash http2Client mkEntitlementProof thParams0) `E.onException` closeHTTP2Client http2Client)
+            `catchError` \e -> liftIO (closeHTTP2Client http2Client) >> throwE e
     _ -> pure thParams0
   logDebug $ "Client negotiated protocol: " <> tshow thVersion
   let c = XFTPClient {http2Client, thParams, transportSession, config}

tests/XFTPClient.hs

@@ -24,10 +24,13 @@ import qualified Simplex.Messaging.Agent as A
 import Simplex.Messaging.Agent.Protocol (UserId)
 import Simplex.Messaging.Crypto.File (CryptoFile)
 import Simplex.Messaging.Protocol (XFTPServer)
-import Simplex.Messaging.Transport.HTTP2 (httpALPN)
+import Simplex.Messaging.Transport (defaultSupportedParams)
+import Simplex.Messaging.Transport.HTTP2 (defaultHTTP2BufferSize, httpALPN)
+import Simplex.Messaging.Transport.HTTP2.Server (runHTTP2Server)
 import Simplex.Messaging.Transport.Server
 import System.Directory (createDirectoryIfMissing)
 import Test.Hspec hiding (fit, it)
+import UnliftIO.STM (TMVar, atomically, newEmptyTMVarIO, putTMVar)
 import Util
 #if defined(dbServerPostgres)
 import Control.Monad (void)
@@ -132,6 +135,19 @@ withXFTPServerCfg cfg@XFTPServerConfig {filesPath} =
     (\started -> createDirectoryIfMissing True filesPath >> runXFTPServerBlocking started cfg)
     (threadDelay 10000)
 
+withStallingXFTPServer :: (HasCallStack, HasTestEnv) => (TMVar () -> TMVar () -> IO a) -> IO a
+withStallingXFTPServer action = do
+  requested <- newEmptyTMVarIO
+  closed <- newEmptyTMVarIO
+  creds <- loadServerCredential xftpCredentials
+  let stall _ _ _ _ _ = atomically (putTMVar requested ()) >> threadDelay maxBound
+  serverBracket
+    (\started -> runHTTP2Server started xftpTestPort defaultHTTP2BufferSize defaultSupportedParams creds Nothing transportConfig Nothing (\_ -> atomically $ putTMVar closed ()) stall)
+    (pure ())
+    (const $ action requested closed)
+  where
+    XFTPServerConfig {xftpCredentials, transportConfig} = testXFTPServerConfig
+
 withXFTPServerThreadOn :: (HasCallStack, HasTestEnv) => (HasCallStack => ThreadId -> IO a) -> IO a
 withXFTPServerThreadOn = withXFTPServerCfg testXFTPServerConfig
 

tests/XFTPServerTests.hs

@@ -5,6 +5,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# OPTIONS_GHC -fno-warn-ambiguous-fields #-}
 
 module XFTPServerTests (xftpServerTests) where
 
@@ -37,7 +38,7 @@ import Simplex.FileTransfer.Description (kb)
 import Simplex.FileTransfer.Protocol (FileInfo (..), XFTPFileId, xftpBlockSize)
 import Simplex.FileTransfer.Server.Env (AFStoreType, XFTPServerConfig (..))
 import Simplex.FileTransfer.Transport (XFTPClientHandshake (..), XFTPClientHello (..), XFTPErrorType (..), XFTPRcvChunkSpec (..), XFTPServerHandshake (..), pattern VersionXFTP)
-import Simplex.Messaging.Client (ProtocolClientError (..))
+import Simplex.Messaging.Client (NetworkConfig (..), NetworkTimeout (..), ProtocolClientError (..), defaultNetworkConfig)
 import qualified Simplex.Messaging.Crypto as C
 import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Encoding (smpDecode, smpEncode)
@@ -80,6 +81,8 @@ xftpServerTests = do
         it "allowed with correct basic auth" $ testFileBasicAuth True (Just "pwd") (Just "pwd") True
         it "allowed with auth on server without auth" $ testFileBasicAuth True Nothing (Just "any") True
       it "should not change content for uploaded and committed files" testFileSkipCommitted
+      it "should close connection when handshake times out" testHandshakeTimeout
+      it "should close connection when cancelled during handshake" testCancelledHandshake
     describe "XFTP SNI and CORS" $ beforeWith (const (pure ())) $ do
       it "should select web certificate when SNI is used" testSNICertSelection
       it "should select XFTP certificate when SNI is not used" testNoSNICertSelection
@@ -388,6 +391,24 @@ testFileSkipCommitted fsType =
         downloadXFTPChunk g c rpKey rId $ XFTPRcvChunkSpec (testPath "received_chunk") chSize digest
         liftIO $ B.readFile (testPath "received_chunk") `shouldReturn` bytes -- new chunk content got ignored
 
+testHandshakeTimeout :: HasTestEnv => AFStoreType -> Expectation
+testHandshakeTimeout _ =
+  withStallingXFTPServer $ \_ closed -> do
+    ts <- getCurrentTime
+    (void <$> getXFTPClient (1, testXFTPServer, Nothing) cfg [] ts (\_ -> pure Nothing) (\_ -> pure ())) `shouldReturn` Left PCEResponseTimeout
+    timeout 1000000 (atomically $ takeTMVar closed) `shouldReturn` Just ()
+  where
+    cfg = testXFTPClientConfig {xftpNetworkConfig = defaultNetworkConfig {tcpConnectTimeout = NetworkTimeout 2000000 2000000}}
+
+testCancelledHandshake :: HasTestEnv => AFStoreType -> Expectation
+testCancelledHandshake _ =
+  withStallingXFTPServer $ \requested closed -> do
+    ts <- getCurrentTime
+    a <- async $ getXFTPClient (1, testXFTPServer, Nothing) testXFTPClientConfig [] ts (\_ -> pure Nothing) (\_ -> pure ())
+    timeout 5000000 (atomically $ takeTMVar requested) `shouldReturn` Just ()
+    cancel a
+    timeout 1000000 (atomically $ takeTMVar closed) `shouldReturn` Just ()
+
 -- SNI and CORS tests
 
 lookupResponseHeader :: B.ByteString -> H2.Response -> Maybe B.ByteString