Safely creating socket

VyacheslavHashov · VyacheslavHashov · commit 1bb7f44b104e · 2017-02-09T10:33:32.000+03:00
diff --git a/src/Database/PostgreSQL/Driver/Connection.hs b/src/Database/PostgreSQL/Driver/Connection.hs
@@ -207,7 +207,7 @@ receiverThread
     -> NotificationHandler
     -> IO ()
 receiverThread msgFilter rawConn dataChan allChan modeRef ntfHandler =
-    receiveLoop []
+    receiveLoop  []
   where
     receiveLoop :: [V.Vector (Maybe B.ByteString)] -> IO ()
     receiveLoop acc = do
diff --git a/src/Database/PostgreSQL/Driver/RawConnection.hs b/src/Database/PostgreSQL/Driver/RawConnection.hs
@@ -2,6 +2,7 @@
 module Database.PostgreSQL.Driver.RawConnection where
 
 import Control.Monad (void)
+import Control.Exception (bracketOnError)
 import Safe (headMay)
 import Data.Monoid ((<>))
 import System.Socket (socket, AddressInfo(..), getAddressInfo, socketAddress,
@@ -38,6 +39,12 @@ createRawConnection settings
         | "/" `B.isPrefixOf` host = unixConnection host
         | otherwise               = tcpConnection
   where
+    createAndConnect Nothing creating = throwAuthErrorInIO AuthInvalidAddress
+    createAndConnect (Just address) creating =
+        bracketOnError creating close $ \s -> do
+            connect s address
+            pure . Right $ constructRawConnection s
+
     unixConnection dirPath = do
         let mAddress = socketAddressUnixPath $ makeUnixPath dirPath
         createAndConnect mAddress (socket :: IO (Socket Unix Stream Unix))
@@ -48,12 +55,6 @@ createRawConnection settings
              :: IO [AddressInfo Inet Stream TCP])
         createAndConnect mAddress (socket :: IO (Socket Inet Stream TCP))
 
-    createAndConnect Nothing creating = throwAuthErrorInIO AuthInvalidAddress
-    createAndConnect (Just address) creating = do
-        s <- creating
-        connect s address
-        pure . Right $ constructRawConnection s
-
     portStr = BS.pack . show $ settingsPort settings
     host    = settingsHost settings
     makeUnixPath dirPath =
