From 6ec0a4590081204b2d7ff857107f7a5f49580a9a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Er=C3=A8be?= <romain.gerard@erebe.eu>
Date: Sat, 21 May 2016 18:05:51 +0200
Subject: [PATCH] Use conduit instead of raw tls package

---
 src/Lib.hs     | 43 +++++++++++--------------------------------
 wstunnel.cabal |  2 +-
 2 files changed, 12 insertions(+), 33 deletions(-)

diff --git a/src/Lib.hs b/src/Lib.hs
index f5123a9..2dd08cb 100644
--- a/src/Lib.hs
+++ b/src/Lib.hs
@@ -31,12 +31,7 @@ import qualified Network.WebSockets            as WS
 import qualified Network.WebSockets.Connection as WS
 import qualified Network.WebSockets.Stream     as WS
 
-import           Network.Connection            (Connection,
-                                                ConnectionParams (..),
-                                                TLSSettings (..), connectTo,
-                                                connectionGetChunk,
-                                                connectionPut,
-                                                initConnectionContext)
+import           Network.Connection            (settingDisableCertificateValidation)
 
 
 instance Hashable N.SockAddr where
@@ -215,34 +210,18 @@ runServer useTLS = if useTLS then runTlsTunnelingServer else runTunnelingServer
 runTlsTunnelingClient :: Proto -> (HostName, PortNumber) -> (HostName, PortNumber) -> (WS.Connection -> IO ()) -> IO ()
 runTlsTunnelingClient proto (wsHost, wsPort) (remoteHost, remotePort) app = do
   putStrLn $ "OPEN tls connection to " <> tshow remoteHost <> ":" <> tshow remotePort
-  context    <- initConnectionContext
-  connection <- connectTo context (connectionParams wsHost (fromIntegral wsPort))
-  stream     <- WS.makeStream (reader connection) (writer connection)
-  WS.runClientWithStream stream wsHost (toPath proto remoteHost remotePort) WS.defaultConnectionOptions [] app
+  let tlsCfg = N.tlsClientConfig (fromIntegral wsPort) (BC.pack wsHost)
+  let tlsSettings = (N.tlsClientTLSSettings tlsCfg) { settingDisableCertificateValidation = True }
+  N.runTLSClient (tlsCfg { N.tlsClientTLSSettings = tlsSettings } )$ \appData ->
+    runApp appData app
+
   putStrLn $ "CLOSE tls connection to " <> tshow remoteHost <> ":" <> tshow remotePort
 
-
-connectionParams :: HostName -> PortNumber -> ConnectionParams
-connectionParams host port = ConnectionParams
-  { connectionHostname = host
-  , connectionPort = port
-  , connectionUseSecure = Just tlsSettings
-  , connectionUseSocks = Nothing
-  }
-
-tlsSettings :: TLSSettings
-tlsSettings = TLSSettingsSimple
-  { settingDisableCertificateValidation = True
-  , settingDisableSession = False
-  , settingUseServerName = False
-  }
-
-reader :: Connection -> IO (Maybe ByteString)
-reader connection = fmap Just (connectionGetChunk connection)
-
-writer :: Connection -> Maybe LByteString -> IO ()
-writer connection = maybe (return ()) (connectionPut connection . toStrict)
-
+  where
+    runApp :: N.AppData -> (WS.Connection -> IO ()) -> IO ()
+    runApp appData app = do
+      stream <- WS.makeStream (Just <$> N.appRead appData) (N.appWrite appData . toStrict . fromJust)
+      WS.runClientWithStream stream wsHost (toPath proto remoteHost remotePort) WS.defaultConnectionOptions [] app
 
 toPath :: Proto -> HostName -> PortNumber -> String
 toPath proto remoteHost remotePort = "/" <> toLower (show proto) <> "/" <> remoteHost <> "/" <> show remotePort
diff --git a/wstunnel.cabal b/wstunnel.cabal
index f0d3840..e178243 100644
--- a/wstunnel.cabal
+++ b/wstunnel.cabal
@@ -23,9 +23,9 @@ library
                      , unordered-containers
                      , network
                      , streaming-commons
-                     , connection
                      , websockets
                      , network-conduit-tls
+                     , connection
 
   default-language:    Haskell2010