1 files changed, 18 insertions, 13 deletions

diff --git a/dht/src/Network/Tox/DHT/Transport.hs b/dht/src/Network/Tox/DHT/Transport.hs
index 5de92916..5f0deea8 100644
--- a/dht/src/Network/Tox/DHT/Transport.hs
+++ b/dht/src/Network/Tox/DHT/Transport.hs
@@ -103,7 +103,7 @@ parseDHTAddr :: (Eq saddr, Show ni) =>
                     (saddr -> STM (Maybe ni))
                     -> (NodeId -> saddr -> Either String ni)
                     -> (ByteString, saddr)
-                    -> IO (Either (DHTMessage Encrypted8,ni) (ByteString,saddr))
+                    -> STM (Either (DHTMessage Encrypted8,ni) (ByteString,saddr))
 parseDHTAddr pendingCookies nodeInfo (msg,saddr)
     | Just (typ,bs) <- B.uncons msg
     , let right = return $ Right (msg,saddr)
@@ -115,9 +115,11 @@ parseDHTAddr pendingCookies nodeInfo (msg,saddr)
     0x04 -> left $ direct nodeInfo bs saddr DHTSendNodes
     0x18 -> left $ direct nodeInfo bs saddr DHTCookieRequest
     0x19 -> do
-        mni <- atomically $ pendingCookies saddr
+        mni <- pendingCookies saddr
         let ni = fromMaybe (noReplyAddr nodeInfo saddr) mni
-        dput XMan $ "Got encrypted cookie! mni="++show mni
+            runio :: IO () -> STM ()
+            runio _ = return () -- TODO: run IO action
+        runio $ dput XMan $ "Got encrypted cookie! mni="++show mni
         left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
     0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo nodeInfo saddr . snd)
     0x21 -> left $ do
@@ -409,13 +411,16 @@ forwardDHTRequests :: TransportCrypto -> (PublicKey -> IO (Maybe ni)) -> DHTTran
 forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
  where
     -- await' :: HandleHi ni a -> STM (IO a)
-    await' pass = awaitMessage dht $ \case
-        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
-                  join $ atomically (await' pass)
-        m -> pass m
+    await' = do
+        (m, io) <- awaitMessage dht
+        return $ case m of
+            Arrival src m@(DHTDHTRequest target payload) | target /= transportPublic crypto
+                -> (,) Discarded $ do
+                      io
+                      mni <- closeLookup target
+                      -- Forward the message if the target is in our close list.
+                      forM_ mni $ \ni -> sendMessage dht ni m
+            _ -> (m,io)
 
 encrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage ((,) Nonce8) -> ni -> IO (DHTMessage Encrypted8, ni)
 encrypt crypto nodeId msg ni = do
@@ -432,7 +437,7 @@ encryptMessage crypto destKey n arg = do
     secret <- lookupSharedSecret crypto (transportSecret crypto) destKey n
     return $ E8 $ ToxCrypto.encrypt secret plain
 
-decrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage Encrypted8 -> ni -> IO (Either String (DHTMessage ((,) Nonce8), ni))
+decrypt :: TransportCrypto -> (ni -> NodeId) -> DHTMessage Encrypted8 -> ni -> STM (Either String (DHTMessage ((,) Nonce8), ni))
 decrypt crypto nodeId msg ni = do
     let decipher n c = Composed $ decryptMessage crypto n . left ((,) $ id2key $ nodeId ni) $ c
     msg' <- sequenceMessage $ transcode decipher msg
@@ -442,11 +447,11 @@ decryptMessage :: Serialize x =>
                   TransportCrypto
                   -> Nonce24
                   -> Either (PublicKey, Encrypted8 x) (Asymm (Encrypted8 x))
-                  -> IO ((Either String ∘ ((,) Nonce8)) x)
+                  -> STM ((Either String ∘ ((,) Nonce8)) x)
 decryptMessage crypto n arg = do
     let (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
         plain8           = Composed . fmap swap . (>>= decodePlain)
-    secret <- lookupSharedSecret crypto (transportSecret crypto) remotekey n
+    secret <- lookupSharedSecretSTM crypto (transportSecret crypto) remotekey n
     return $ plain8 $ ToxCrypto.decrypt secret e
 
 sequenceMessage :: Applicative m => DHTMessage (m ∘ f) -> m (DHTMessage f)