2 files changed, 10 insertions, 10 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 ()
diff --git a/dht/src/Network/Tox/DHT/Transport.hs b/dht/src/Network/Tox/DHT/Transport.hs
index 7475b3b1..7414343d 100644
--- a/dht/src/Network/Tox/DHT/Transport.hs
+++ b/dht/src/Network/Tox/DHT/Transport.hs
@@ -59,7 +59,7 @@ import GHC.Generics
 import Network.Socket
 
 type DHTTransport = Transport String NodeInfo (DHTMessage Encrypted8)
-type HandleHi a = Maybe (Either String (DHTMessage Encrypted8, NodeInfo)) -> IO a
+type HandleHi a = Arrival String NodeInfo (DHTMessage Encrypted8) -> IO a
 
 
 data DHTMessage (f :: * -> *)
@@ -399,13 +399,13 @@ instance Serialize CookieRequest where
 forwardDHTRequests :: TransportCrypto -> (PublicKey -> IO (Maybe NodeInfo)) -> DHTTransport -> DHTTransport
 forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
  where
-    await' :: HandleHi a -> IO a
+    await' :: HandleHi a -> STM (IO a)
     await' pass = awaitMessage dht $ \case
-        Just (Right (m@(DHTDHTRequest target payload),src)) | target /= transportPublic crypto
+        Arrival src m@(DHTDHTRequest target payload) | target /= transportPublic crypto
             -> do mni <- closeLookup target
                   -- Forward the message if the target is in our close list.
                   forM_ mni $ \ni -> sendMessage dht ni m
-                  await' pass
+                  join $ atomically (await' pass)
         m -> pass m
 
 encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> IO (DHTMessage Encrypted8, NodeInfo)