STM-based awaitMessage.

1 files changed, 6 insertions, 6 deletions

diff --git a/dht/src/Network/Tox/AggregateSession.hs b/dht/src/Network/Tox/AggregateSession.hs
index 8c728660..999c7399 100644
--- a/dht/src/Network/Tox/AggregateSession.hs
+++ b/dht/src/Network/Tox/AggregateSession.hs
@@ -188,21 +188,21 @@ forkSession c s setStatus = forkIO $ do
             now <- getPOSIXTime
             atomically $ modifyTVar' q $ PSQ.insert (fromEnum DoTimeout) (now + 15)
 
-        onPacket body loop Nothing          = return ()
-        onPacket body loop (Just (Left e))  = inPrint e >> loop
-        onPacket body loop (Just (Right x)) = body loop x
+        onPacket body loop Terminated     = return ()
+        onPacket body loop (ParseError e) = inPrint e >> loop
+        onPacket body loop (Arrival _ x)  = body loop x
 
-        awaitPacket body = fix $ awaitMessage (sTransport s) . onPacket body
+        awaitPacket body = fix $ join . atomically . awaitMessage (sTransport s) . onPacket body
 
     atomically $ setStatus $ InProgress AwaitingSessionPacket
-    awaitPacket $ \_ (online,()) -> do
+    awaitPacket $ \_ online -> do
         when (msgID online /= M ONLINE) $ do
             inPrint $ "Unexpected initial packet: " ++ show (msgID online)
         atomically $ do setStatus Established
                         sendPacket online
         bump
         beacon <- forkIO $ keepAlive s q
-        awaitPacket $ \awaitNext (x,()) -> do
+        awaitPacket $ \awaitNext x -> do
             bump
             case msgID x of
                 M ALIVE      -> return ()