From 1df146c702cab5a4ef963ef9f33b6154af1a5d5a Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Thu, 16 Dec 2021 07:30:16 +0000
Subject: [PATCH] make WebSockets use TLS Context (#227)

---
 src/Simplex/Messaging/Transport.hs            | 32 +++++++++++--------
 src/Simplex/Messaging/Transport/WebSockets.hs | 19 ++++++-----
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/src/Simplex/Messaging/Transport.hs b/src/Simplex/Messaging/Transport.hs
index 95ded7962..101bf5a6b 100644
--- a/src/Simplex/Messaging/Transport.hs
+++ b/src/Simplex/Messaging/Transport.hs
@@ -102,10 +102,10 @@ class Transport c where
   transportName :: TProxy c -> String
 
   -- | Upgrade client TLS context to connection (used in the server)
-  getServerConnection :: TLS -> IO c
+  getServerConnection :: T.Context -> IO c
 
   -- | Upgrade server TLS context to connection (used in the client)
-  getClientConnection :: TLS -> IO c
+  getClientConnection :: T.Context -> IO c
 
   -- | Close connection
   closeConnection :: c -> IO ()
@@ -151,7 +151,8 @@ runTransportServer started port serverParams server = do
     acceptConnection :: Transport c => Socket -> IO c
     acceptConnection sock = do
       (newSock, _) <- accept sock
-      connectTLS "server" getServerConnection serverParams newSock
+      ctx <- connectTLS "server" serverParams newSock
+      getServerConnection ctx
 
 startTCPServer :: TMVar Bool -> ServiceName -> IO Socket
 startTCPServer started port = withSocketsDo $ resolve >>= open >>= setStarted
@@ -194,7 +195,8 @@ startTCPClient host port = withSocketsDo $ resolve >>= tryOpen err
     open addr = do
       sock <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
       connect sock $ addrAddress addr
-      connectTLS "client" getClientConnection clientParams sock
+      ctx <- connectTLS "client" clientParams sock
+      getClientConnection ctx
 
 loadTLSServerParams :: FilePath -> FilePath -> IO T.ServerParams
 loadTLSServerParams certificateFile privateKeyFile =
@@ -218,14 +220,18 @@ loadTLSServerParams certificateFile privateKeyFile =
 
 data TLS = TLS {tlsContext :: T.Context, buffer :: TVar ByteString, getLock :: TMVar ()}
 
-connectTLS :: (T.TLSParams p) => String -> (TLS -> IO c) -> p -> Socket -> IO c
-connectTLS party getPartyConnection params sock =
-  E.bracketOnError (T.contextNew sock params) closeTLS $ \tlsContext -> do
-    T.handshake tlsContext
-    buffer <- newTVarIO ""
-    getLock <- newTMVarIO ()
-    getPartyConnection TLS {tlsContext, buffer, getLock}
+connectTLS :: T.TLSParams p => String -> p -> Socket -> IO T.Context
+connectTLS party params sock =
+  E.bracketOnError (T.contextNew sock params) closeTLS $ \ctx -> do
+    T.handshake ctx
       `E.catch` \(e :: E.SomeException) -> putStrLn (party <> " exception: " <> show e) >> E.throwIO e
+    pure ctx
+
+newTLS :: T.Context -> IO TLS
+newTLS tlsContext = do
+  buffer <- newTVarIO ""
+  getLock <- newTMVarIO ()
+  pure TLS {tlsContext, buffer, getLock}
 
 closeTLS :: T.Context -> IO ()
 closeTLS ctx =
@@ -252,8 +258,8 @@ supportedParameters =
 
 instance Transport TLS where
   transportName _ = "TLS 1.3"
-  getServerConnection = pure
-  getClientConnection = pure
+  getServerConnection = newTLS
+  getClientConnection = newTLS
   closeConnection tls = closeTLS $ tlsContext tls
 
   cGet :: TLS -> Int -> IO ByteString
diff --git a/src/Simplex/Messaging/Transport/WebSockets.hs b/src/Simplex/Messaging/Transport/WebSockets.hs
index 4e0ba0ffb..4aaea44df 100644
--- a/src/Simplex/Messaging/Transport/WebSockets.hs
+++ b/src/Simplex/Messaging/Transport/WebSockets.hs
@@ -1,6 +1,5 @@
 {-# LANGUAGE InstanceSigs #-}
 {-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE NamedFieldPuns #-}
 
 module Simplex.Messaging.Transport.WebSockets (WS (..)) where
 
@@ -12,7 +11,7 @@ import qualified Network.TLS as T
 import Network.WebSockets
 import Network.WebSockets.Stream (Stream)
 import qualified Network.WebSockets.Stream as S
-import Simplex.Messaging.Transport (TLS (..), TProxy, Transport (..), TransportError (..), closeTLS, trimCR)
+import Simplex.Messaging.Transport (TProxy, Transport (..), TransportError (..), closeTLS, trimCR)
 
 data WS = WS {wsStream :: Stream, wsConnection :: Connection}
 
@@ -28,17 +27,17 @@ instance Transport WS where
   transportName :: TProxy WS -> String
   transportName _ = "WebSockets"
 
-  getServerConnection :: TLS -> IO WS
-  getServerConnection TLS {tlsContext} = do
-    s <- websocketsStream tlsContext
+  getServerConnection :: T.Context -> IO WS
+  getServerConnection ctx = do
+    s <- makeTLSContextStream ctx
     WS s <$> acceptClientRequest s
     where
       acceptClientRequest :: Stream -> IO Connection
       acceptClientRequest s = makePendingConnectionFromStream s websocketsOpts >>= acceptRequest
 
-  getClientConnection :: TLS -> IO WS
-  getClientConnection TLS {tlsContext} = do
-    s <- websocketsStream tlsContext
+  getClientConnection :: T.Context -> IO WS
+  getClientConnection ctx = do
+    s <- makeTLSContextStream ctx
     WS s <$> sendClientRequest s
     where
       sendClientRequest :: Stream -> IO Connection
@@ -64,8 +63,8 @@ instance Transport WS where
       then E.throwIO TEBadBlock
       else pure $ B.init s
 
-websocketsStream :: T.Context -> IO S.Stream
-websocketsStream tlsContext =
+makeTLSContextStream :: T.Context -> IO S.Stream
+makeTLSContextStream tlsContext =
   S.makeStream readStream writeStream
   where
     readStream :: IO (Maybe ByteString)