close stream after failed connection attempt (including exceptions)

source/Network/Xmpp/Concurrent.hs

@@ -170,6 +170,39 @@ newSession stream config realm mbSasl = runErrorT $ do
                                                  onConnectionClosed config sess }
     return sess
 
+connectStream :: HostName
+              -> SessionConfiguration
+              -> Maybe (ConnectionState -> [SaslHandler], Maybe Text)
+              -> IO (Either XmppFailure Stream)
+connectStream realm config mbSasl = do
+    Ex.bracketOnError (openStream realm (sessionStreamConfiguration config))
+                      (\s -> case s of
+                                  Left _ -> return ()
+                                  Right stream -> closeStreams stream)
+
+        (\stream' -> case stream' of
+              Left e -> return $ Left e
+              Right stream -> do
+                  res <- runErrorT $ do
+                      ErrorT $ tls stream
+                      cs <- liftIO $ withStream (gets streamConnectionState)
+                                                stream
+                      mbAuthError <- case mbSasl of
+                          Nothing -> return Nothing
+                          Just (handlers, resource) -> ErrorT $ auth (handlers cs)
+                                                                resource stream
+                      case mbAuthError of
+                          Nothing -> return ()
+                          Just e  -> throwError $ XmppAuthFailure e
+                      return stream
+                  case res of
+                      Left e -> do
+                          debugM "Pontarius.Xmpp" "Closing stream after error"
+                          closeStreams stream
+                          return (Left e)
+                      Right r -> return $ Right r
+                  )
+
 -- | Creates a 'Session' object by setting up a connection with an XMPP server.
 --
 -- Will connect to the specified host with the provided configuration. If the
@@ -182,15 +215,7 @@ session :: HostName                          -- ^ The hostname / realm
         -> SessionConfiguration              -- ^ configuration details
         -> IO (Either XmppFailure Session)
 session realm mbSasl config = runErrorT $ do
-    stream <- ErrorT $ openStream realm (sessionStreamConfiguration config)
+    stream <- ErrorT $ connectStream realm config mbSasl
-    ErrorT $ tls stream
-    cs <- liftIO $ withStream (gets streamConnectionState) stream
-    mbAuthError <- case mbSasl of
-        Nothing -> return Nothing
-        Just (handlers, resource) -> ErrorT $ auth (handlers cs) resource stream
-    case mbAuthError of
-        Nothing -> return ()
-        Just e  -> throwError $ XmppAuthFailure e
     ses <- ErrorT $ newSession stream config realm mbSasl
     liftIO $ when (enableRoster config) $ initRoster ses
     return ses
@@ -202,24 +227,10 @@ reconnectNow :: Session -- ^ session to reconnect
 reconnectNow sess@Session{conf = config, reconnectWait = rw} = do
     debugM "Pontarius.Xmpp" "reconnecting"
     res <- flip withConnection sess $ \oldStream -> do
-        s <- runErrorT $ do
+        debugM "Pontarius.Xmpp" "reconnect: closing stream"
-            liftIO $ debugM "Pontarius.Xmpp" "reconnect: closing stream"
+        closeStreams oldStream
-            _ <- liftIO $ closeStreams oldStream
+        debugM "Pontarius.Xmpp" "reconnect: opening stream"
-            liftIO $ debugM "Pontarius.Xmpp" "reconnect: opening stream"
+        s <- connectStream (sRealm sess) config (sSaslCredentials sess)
-            stream <- ErrorT $ openStream (sRealm sess)
-                                          (sessionStreamConfiguration config)
-            liftIO $ debugM "Pontarius.Xmpp" "reconnect: tls"
-            ErrorT $ tls stream
-            liftIO $ debugM "Pontarius.Xmpp" "reconnect: auth"
-            cs <- liftIO $ withStream (gets streamConnectionState) stream
-            mbAuthError <- case sSaslCredentials sess of
-                Nothing -> return Nothing
-                Just (handlers, resource) -> ErrorT $ auth (handlers cs)
-                                                           resource stream
-            case mbAuthError of
-                Nothing -> return ()
-                Just e  -> throwError $ XmppAuthFailure e
-            return stream
         case s of
             Left  e -> do
                 errorM "Pontarius.Xmpp" $ "reconnect failed"  ++ show e

tests/Tests.hs

@@ -218,8 +218,8 @@ connectionClosedTest = do
                        }
   sendPresence presenceOnline context
   forkIO $ do
-      threadDelay 1000000
+      threadDelay 5000000
-      endSession context
+      closeConnection context
       debug' "done"
   forever $ threadDelay 1000000
   return ()