From 33fe3554b8cdfca626ee7fcdfb9b222ee8cc7f3c Mon Sep 17 00:00:00 2001
From: "Evgeny @ SimpleX Chat"
 <259188159+evgeny-simplex@users.noreply.github.com>
Date: Wed, 21 Oct 2026 14:11:23 +0000
Subject: [PATCH] client: kill transport thread on interrupt

---
 src/Simplex/Messaging/Client.hs                 |  6 +++---
 src/Simplex/Messaging/Transport/HTTP2/Client.hs |  8 ++++----
 tests/SMPProxyTests.hs                          | 17 ++++++++++++++++-
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 9989a96eb..bf2090134 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -611,16 +611,16 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
           }
 
     runClient :: (ServiceName, ATransport 'TClient) -> TransportHost -> PClient v err msg -> IO (Either (ProtocolClientError err) (ProtocolClient v err msg))
-    runClient (port', ATransport t) useHost c = do
+    runClient (port', ATransport t) useHost c = E.mask $ \restore -> do
       cVar <- newEmptyTMVarIO
       let tcConfig = (transportClientConfig networkConfig nm useHost useSNI useALPN) {clientCredentials = serviceCreds <$> serviceCredentials}
           socksCreds = clientSocksCredentials networkConfig proxySessTs transportSession
       tId <-
-        runTransportClient tcConfig socksCreds useHost port' (Just $ keyHash srv) (client t c cVar)
+        restore (runTransportClient tcConfig socksCreds useHost port' (Just $ keyHash srv) (client t c cVar))
           `forkFinally` \r ->
             let err = either toNetworkError (const NEFailedError) r
              in void $ atomically $ tryPutTMVar cVar $ Left $ PCENetworkError err
-      c_ <- netTimeoutInt tcpConnectTimeout nm `timeout` atomically (takeTMVar cVar)
+      c_ <- restore (netTimeoutInt tcpConnectTimeout nm `timeout` atomically (takeTMVar cVar)) `E.onException` killThread tId
       case c_ of
         Just (Right c') -> mkWeakThreadId tId >>= \tId' -> pure $ Right c' {action = Just tId'}
         Just (Left e) -> pure $ Left e
diff --git a/src/Simplex/Messaging/Transport/HTTP2/Client.hs b/src/Simplex/Messaging/Transport/HTTP2/Client.hs
index ca0714225..170e8740f 100644
--- a/src/Simplex/Messaging/Transport/HTTP2/Client.hs
+++ b/src/Simplex/Messaging/Transport/HTTP2/Client.hs
@@ -140,13 +140,13 @@ getVerifiedHTTP2ClientWith config host port disconnected setup =
       pure HClient {connected, disconnected, host, port, config, reqQ}
 
     runClient :: HClient -> IO (Either HTTP2ClientError HTTP2Client)
-    runClient c = do
+    runClient c = E.mask $ \restore -> do
       cVar <- newEmptyTMVarIO
       action <-
-        async $ setup (client c cVar) `E.catch` \e -> do
-          atomically $ putTMVar cVar $ Left $ HCNetworkError $ toNetworkError e
+        async $ restore (setup (client c cVar)) `E.catch` \e -> do
+          void $ atomically $ tryPutTMVar cVar $ Left $ HCNetworkError $ toNetworkError e
           E.throwIO e
-      c_ <- connTimeout config `timeout` atomically (takeTMVar cVar)
+      c_ <- restore (connTimeout config `timeout` atomically (takeTMVar cVar)) `E.onException` cancel action
       case c_ of
         Just (Right c') -> pure $ Right c' {action = Just action}
         Just (Left e) -> pure $ Left e
diff --git a/tests/SMPProxyTests.hs b/tests/SMPProxyTests.hs
index 6a6e22069..774f318d0 100644
--- a/tests/SMPProxyTests.hs
+++ b/tests/SMPProxyTests.hs
@@ -21,9 +21,11 @@ import Control.Logger.Simple
 import Control.Monad (forM, forM_, forever, replicateM_)
 import Control.Monad.Trans.Except (ExceptT, runExceptT)
 import Data.ByteString.Char8 (ByteString)
+import qualified Data.ByteString.Char8 as B
 import Data.List.NonEmpty (NonEmpty)
 import qualified Data.List.NonEmpty as L
 import Data.Time.Clock (getCurrentTime)
+import Network.Socket.ByteString (recv)
 import SMPAgentClient
 import SMPClient
 import ServerTests (decryptMsgV3, responseTime, sendRecv)
@@ -41,7 +43,8 @@ import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Server.Env.STM (AStoreType (..), ServerConfig (..))
 import Simplex.Messaging.Server.MsgStore.Types (SQSType (..))
 import Simplex.Messaging.Transport
-import Simplex.Messaging.Util (bshow, tshow)
+import Simplex.Messaging.Transport.Server (runLocalTCPServer)
+import Simplex.Messaging.Util (bshow, tshow, unlessM)
 import Simplex.Messaging.Version (mkVersionRange)
 import System.FilePath (splitExtensions)
 import System.Random (randomRIO)
@@ -59,6 +62,8 @@ smpProxyTests = do
       testProxyRecoversWithoutDisconnect
     it "reconnects to relay after sender disconnects mid-connection" $ \_ ->
       testProxyReconnectAfterRelayRestart
+    it "closes connection to relay after sender disconnects mid-connection" $ \_ ->
+      testProxyClosesRelayConnectionAfterDisconnect
   describe "agent client reconnection" $ do
     it "reconnects after a connect is cancelled mid-flight" $ \_ ->
       testAgentClientReconnectAfterCancel
@@ -558,6 +563,16 @@ testProxyReconnectAfterRelayRestart =
       race_ (threadDelay 1000000) requestRelaySession
     requireProxyReconnect
 
+testProxyClosesRelayConnectionAfterDisconnect :: HasTestEnv => IO ()
+testProxyClosesRelayConnectionAfterDisconnect =
+  withSmpServerConfigOn (transport @TLS) proxyCfgShortTimeout testPort $ \_ -> do
+    closed <- newEmptyTMVarIO
+    let waitClosed sock = unlessM (B.null <$> recv sock 4096) $ waitClosed sock
+        stallingRelay started = runLocalTCPServer started testPort2 $ \sock -> waitClosed sock >> atomically (putTMVar closed ())
+    serverBracket stallingRelay (pure ()) $ \_ -> do
+      race_ (threadDelay 1000000) requestRelaySession
+      timeout 1000000 (atomically $ takeTMVar closed) `shouldReturn` Just ()
+
 -- Bug B (same root cause as the proxy, in the messaging agent): getSMPServerClient inserts an
 -- empty SessionVar into smpClients, then connects inside newProtocolClient's tryAllErrors, which
 -- rethrows async exceptions. If the connecting thread is cancelled mid-connect, putTMVar is