1 files changed, 15 insertions, 12 deletions

diff --git a/dht/src/Network/Tox/AggregateSession.hs b/dht/src/Network/Tox/AggregateSession.hs
index 9a784291..d1f42e91 100644
--- a/dht/src/Network/Tox/AggregateSession.hs
+++ b/dht/src/Network/Tox/AggregateSession.hs
@@ -24,6 +24,7 @@ module Network.Tox.AggregateSession
 
 import Control.Concurrent.STM
 import Control.Concurrent.STM.TMChan
+import Control.Exception
 import Control.Monad
 import Data.Dependent.Sum
 import Data.Function
@@ -114,7 +115,7 @@ data KeepAliveEvents = DoTimeout        -- ^ A session timed-out, close it.
                      | DoRequestMissing -- ^ Detect and request lost packets.
                      deriving Enum
 
--- | This call loops until the provided sesison is closed or times out.  It
+-- | This call loops until the provided session is closed or times out.  It
 -- monitors the provided (non-empty) priority queue for scheduled tasks (see
 -- 'KeepAliveEvents') to perform for the connection.
 keepAlive :: Session -> TVar (PSQ POSIXTime) -> IO ()
@@ -142,23 +143,23 @@ keepAlive s q = do
                     now <- getPOSIXTime
                     atomically $ modifyTVar' q $ PSQ.insert (fromEnum DoTimeout) now
 
-        re tm again e io = do
+        re tm e io = do
             io
             atomically $ modifyTVar' q $ PSQ.insert (fromEnum e) tm
-            again
 
         doEvent again now e = case e of
             DoTimeout        -> do dput XNetCrypto $ "TIMEOUT: " ++ show (sSessionID s)
                                    sClose s
-            DoAlive          -> re (now + 10) again e doAlive
-            DoRequestMissing -> re (now + 5)  again e doRequestMissing -- tox-core does this at 1 second intervals
+            DoAlive          -> re (now + 10) e doAlive >> again
+            DoRequestMissing -> re (now + 5{- toxcore uses 1sec -}) e doRequestMissing >> again
 
     fix $ \again -> do
 
     now <- getPOSIXTime
     join $ atomically $ do
         PSQ.findMin <$> readTVar q >>= \case
-          Nothing           -> error "keepAlive: unexpected empty PSQ."
+          Nothing           -> return $ do dput XUnexpected "keepAlive: unexpected empty PSQ."
+                                           sClose s
           Just ( k :-> tm ) ->
             return $ if now < tm then threadDelay (toMicroseconds $ tm - now) >> again
                                  else doEvent again now (toEnum k)
@@ -202,7 +203,7 @@ forkSession c s setStatus = forkIO $ do
         atomically $ do setStatus Established
                         sendPacket online
         bump
-        beacon <- forkIO $ keepAlive s q
+        beacon <- forkIO $ keepAlive s q `finally` sClose s
         awaitPacket $ \awaitNext x -> do
             bump
             case msgID x of
@@ -223,7 +224,7 @@ forkSession c s setStatus = forkIO $ do
 -- one active session).
 addSession :: AggregateSession -> Session -> IO AddResult
 addSession c s = do
-    (result,mcon,replaced) <- atomically $ do
+    (result,mcon,rejected) <- atomically $ do
         let them = sTheirUserKey s
             me   = toPublic $ sOurKey s
         compat <- checkCompatible me them c
@@ -232,7 +233,7 @@ addSession c s = do
                         Just True    -> AddedSession
                         Just False   -> RejectedSession
         case result of
-            RejectedSession -> return (result,Nothing,Nothing)
+            RejectedSession -> return (result,Nothing,Just s)
             _               -> do
                                 statvar <- newTVar Dormant
                                 imap <- readTVar (contactSession c)
@@ -240,9 +241,9 @@ addSession c s = do
                                     s0 = IntMap.lookup (sSessionID s) imap
                                     imap' = IntMap.insert (sSessionID s) con imap
                                 writeTVar (contactSession c) imap'
-                                return (result,Just con,s0)
+                                return (result,Just con,singleSession <$> s0)
 
-    mapM_ (sClose . singleSession) replaced
+    mapM_ sClose rejected
     forM_ (mcon :: Maybe SingleCon) $ \con ->
         forkSession c s $ \progress -> do
             status0 <- aggregateStatus c
@@ -313,7 +314,9 @@ closeAll :: AggregateSession -> IO ()
 closeAll c = join $ atomically $ do
     imap <- readTVar (contactSession c)
     closeTMChan (contactChannel c)
-    return $ forM_ (IntMap.keys imap) $ \sid -> delSession c sid
+    return $ forM_ (IntMap.toList imap) $ \(sid,SingleCon s _) -> do
+        sClose s
+        delSession c sid
 
 -- | Query the current status of the aggregate, there are three possible
 -- values: