STM-based awaitMessage.

1 files changed, 4 insertions, 4 deletions

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)