Exception handling for TCP-relay client connect.

1 files changed, 6 insertions, 3 deletions

diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
index c0bdcd3c..154e9145 100644
--- a/src/Network/QueryResponse/TCP.hs
+++ b/src/Network/QueryResponse/TCP.hs
@@ -67,9 +67,12 @@ acquireConnection mvar tcpcache stream addr = do
     case MM.lookup' (TCPAddress $ streamAddr stream addr) cache of
         Nothing -> do
             proto <- getProtocolNumber "tcp"
-            sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
-            connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
-            h <- socketToHandle sock ReadWriteMode
+            mh <- catchIOError (do sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
+                                   connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
+                                   h <- socketToHandle sock ReadWriteMode
+                                   return $ Just h)
+                               $ \e -> return Nothing
+            fmap join $ forM mh $ \h -> do
             st <- streamHello stream addr h
             t <- getPOSIXTime
             rthread <- forkIO $ fix $ \loop -> do