From c501f4f9ccdd48807a5153697ea1827129841158 Mon Sep 17 00:00:00 2001
From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com>
Date: Fri, 17 Nov 2023 12:37:32 +0200
Subject: [PATCH] remote: fix deadlocked client when server rejects its cert
 (#897)

* remote: detect tls errors sooner to prevent deadlocks

* remove redundant error checking

* cleanup
---
 src/Simplex/Messaging/Transport/Buffer.hs |  9 +++++++++
 src/Simplex/RemoteControl/Client.hs       | 18 ++++++++++--------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/src/Simplex/Messaging/Transport/Buffer.hs b/src/Simplex/Messaging/Transport/Buffer.hs
index 251471679..6de9326f8 100644
--- a/src/Simplex/Messaging/Transport/Buffer.hs
+++ b/src/Simplex/Messaging/Transport/Buffer.hs
@@ -6,6 +6,7 @@ module Simplex.Messaging.Transport.Buffer where
 
 import Control.Concurrent.STM
 import qualified Control.Exception as E
+import Control.Monad (forM_)
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import GHC.IO.Exception (IOErrorType (..), IOException (..), ioException)
@@ -28,6 +29,14 @@ withBufferLock TBuffer {getLock} =
     (atomically $ takeTMVar getLock)
     (atomically $ putTMVar getLock ())
 
+-- | Attempt to read some bytes, appending it to the existing buffer
+peekBuffered :: TBuffer -> Int -> IO ByteString -> IO (ByteString, Maybe ByteString)
+peekBuffered tb@TBuffer {buffer} t getChunk = withBufferLock tb $ do
+  old <- readTVarIO buffer
+  next_ <- timeout t getChunk
+  forM_ next_ $ \next -> atomically $ writeTVar buffer $! old <> next
+  pure (old, next_)
+
 getBuffered :: TBuffer -> Int -> Maybe Int -> IO ByteString -> IO ByteString
 getBuffered tb@TBuffer {buffer} n t_ getChunk = withBufferLock tb $ do
   b <- readChunks True =<< readTVarIO buffer
diff --git a/src/Simplex/RemoteControl/Client.hs b/src/Simplex/RemoteControl/Client.hs
index 2f9050288..efbda4998 100644
--- a/src/Simplex/RemoteControl/Client.hs
+++ b/src/Simplex/RemoteControl/Client.hs
@@ -54,7 +54,8 @@ import qualified Simplex.Messaging.Crypto.Lazy as LC
 import Simplex.Messaging.Crypto.SNTRUP761
 import Simplex.Messaging.Crypto.SNTRUP761.Bindings
 import Simplex.Messaging.Encoding
-import Simplex.Messaging.Transport (TLS (tlsUniq), cGet, cPut)
+import Simplex.Messaging.Transport (TLS (..), cGet, cPut)
+import Simplex.Messaging.Transport.Buffer (peekBuffered)
 import Simplex.Messaging.Transport.Client (TransportClientConfig (..), TransportHost, defaultTransportClientConfig, runTransportClient)
 import Simplex.Messaging.Transport.Credentials (genCredentials, tlsCredentials)
 import Simplex.Messaging.Util
@@ -273,13 +274,14 @@ connectRCCtrl_ drg pairing'@RCCtrlPairing {caKey, caCert} inv@RCInvitation {ca,
           TLS.Credentials (creds : _) -> pure $ Just creds
           _ -> throwError $ RCEInternal "genTLSCredentials must generate credentials"
       let clientConfig = defaultTransportClientConfig {clientCredentials}
-      liftIO . runTransportClient clientConfig Nothing host (show port) (Just ca) $ \tls ->
-        void . runExceptT $ do
-          logDebug "Got TLS connection"
-          r' <- newEmptyTMVarIO
-          whenM (atomically $ tryPutTMVar r $ Right (tlsUniq tls, tls, r')) $ do
-            logDebug "Waiting for session confirmation"
-            whenM (atomically $ readTMVar confirmSession) (runSession tls r') `putRCError` r'
+      runTransportClient clientConfig Nothing host (show port) (Just ca) $ \tls@TLS {tlsBuffer, tlsContext} -> do
+        -- pump socket to detect connection problems
+        liftIO $ peekBuffered tlsBuffer 100000 (TLS.recvData tlsContext) >>= logDebug . tshow -- should normally be ("", Nothing) here
+        logDebug "Got TLS connection"
+        r' <- newEmptyTMVarIO
+        whenM (atomically $ tryPutTMVar r $ Right (tlsUniq tls, tls, r')) $ do
+          logDebug "Waiting for session confirmation"
+          whenM (atomically $ readTMVar confirmSession) $ runSession tls r' `putRCError` r'
       where
         runSession tls r' = do
           (sharedKey, kemPrivKey, hostEncHello) <- prepareHostHello drg pairing' inv hostAppInfo