5 files changed, 48 insertions, 39 deletions

diff --git a/dht/src/Network/Tox/AggregateSession.hs b/dht/src/Network/Tox/AggregateSession.hs
index feb634f0..33b1fafb 100644
--- a/dht/src/Network/Tox/AggregateSession.hs
+++ b/dht/src/Network/Tox/AggregateSession.hs
@@ -196,7 +196,10 @@ forkSession c s setStatus = forkIO $ do
         onPacket body loop (ParseError e) = inPrint e >> loop
         onPacket body loop (Arrival _ x)  = body loop x
 
-        awaitPacket body = fix $ join . atomically . awaitMessage (sTransport s) . onPacket body
+        awaitPacket body = fix $ \loop -> do
+            (m,io) <- atomically $ awaitMessage (sTransport s)
+            io
+            onPacket body loop m
 
     atomically $ setStatus $ InProgress AwaitingSessionPacket
     awaitPacket $ \_ online -> do
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)
diff --git a/dht/src/Network/Tox/Onion/Routes.hs b/dht/src/Network/Tox/Onion/Routes.hs
index 46ded48d..93e9bfcd 100644
--- a/dht/src/Network/Tox/Onion/Routes.hs
+++ b/dht/src/Network/Tox/Onion/Routes.hs
@@ -539,16 +539,16 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
         Nothing -> routeLogger or $ "ONION Failed RouteId " ++ show rid
 
 
-lookupSender :: OnionRouter -> SockAddr -> Nonce8 -> IO (Maybe (OnionDestination RouteId))
-lookupSender or = lookupSender' (pendingQueries or) (routeLog or)
+lookupSender :: OnionRouter -> SockAddr -> Nonce8 -> STM (Maybe (OnionDestination RouteId))
+lookupSender or saddr n8 = lookupSender' (pendingQueries or) (routeLog or) saddr n8
 
 lookupSender' :: TVar (Word64Map PendingQuery)
                        -> TChan String
                        -> SockAddr
                        -> Nonce8
-                       -> IO (Maybe (OnionDestination RouteId))
+                       -> STM (Maybe (OnionDestination RouteId))
 lookupSender' pending log saddr (Nonce8 w8) = do
-    result <- atomically $ do
+    result <- do
         ks <- readTVar pending
         let r = W64.lookup w8 ks
         writeTChan log $ "ONION lookupSender " ++ unwords [show w8, "->", show r]
diff --git a/dht/src/Network/Tox/Session.hs b/dht/src/Network/Tox/Session.hs
index d34dfc7a..53d63287 100644
--- a/dht/src/Network/Tox/Session.hs
+++ b/dht/src/Network/Tox/Session.hs
@@ -106,14 +106,13 @@ sClose s = do
 -- negotiated.  It always returns Nothing which makes it convenient to use with
 -- 'Network.QueryResponse.addHandler'.
 handshakeH :: SessionParams
-              -> Multi.SessionAddress
-              -> Handshake Encrypted
-              -> IO (Maybe a)
-handshakeH sp saddr handshake = do
+              -> Arrival err Multi.SessionAddress (Handshake Encrypted)
+              -> STM (Arrival err Multi.SessionAddress (Handshake Encrypted), IO ())
+handshakeH sp (Arrival saddr handshake) = return $ (,) Discarded $ do
     decryptHandshake (spCrypto sp) handshake
         >>= either (\err -> return ())
                    (uncurry $ plainHandshakeH sp saddr)
-    return Nothing
+handshakeH _ m = return (m, return ())
 
 
 plainHandshakeH :: SessionParams
diff --git a/dht/src/Network/Tox/TCP.hs b/dht/src/Network/Tox/TCP.hs
index 1da302b6..626d4714 100644
--- a/dht/src/Network/Tox/TCP.hs
+++ b/dht/src/Network/Tox/TCP.hs
@@ -160,22 +160,22 @@ tcpStream crypto mkst = StreamHandshake
     , streamAddr = nodeAddr
     }
 
-newtype SessionData = SessionData (MVar (IntMap.IntMap NodeId))
+newtype SessionData = SessionData (TMVar (IntMap.IntMap NodeId))
 
 newSessionData :: NodeInfo -> IO SessionData
-newSessionData _ = SessionData <$> newMVar IntMap.empty
+newSessionData _ = atomically $ SessionData <$> newTMVar IntMap.empty
 
-getRelayedRemote :: SessionData -> ConId -> IO NodeId
+getRelayedRemote :: SessionData -> ConId -> STM NodeId
 getRelayedRemote (SessionData keymapVar) (ConId i) = do
-    keymap <- takeMVar keymapVar
+    keymap <- takeTMVar keymapVar
     let k = fromMaybe UDP.zeroID $ IntMap.lookup (fromIntegral i) keymap
-    putMVar keymapVar keymap
+    putTMVar keymapVar keymap
     return k
 
-setRelayedRemote :: SessionData -> ConId -> NodeId -> IO ()
+setRelayedRemote :: SessionData -> ConId -> NodeId -> STM ()
 setRelayedRemote (SessionData keymapVar) (ConId conid) nid = do
-    keymap <- takeMVar keymapVar
-    putMVar keymapVar $ IntMap.insert (fromIntegral conid) nid keymap
+    keymap <- takeTMVar keymapVar
+    putTMVar keymapVar $ IntMap.insert (fromIntegral conid) nid keymap
 
 toxTCP :: TransportCrypto -> IO ( TCPCache (SessionProtocol (SessionData,RelayPacket) RelayPacket)
                                 , TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket) )
@@ -367,7 +367,7 @@ type RelayCache = TCPCache (SessionProtocol (SessionData,RelayPacket) RelayPacke
 newClient :: TransportCrypto
                -> ((Nonce8 -> QR.Result (Bool,RelayPacket) -> IO ()) -> a)      -- ^ store mvar for relay query
                -> (a -> Nonce8 -> RelayPacket -> IO void)                       -- ^ load mvar for relay query
-               -> (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination RouteId))) -- ^ lookup sender of onion query
+               -> (SockAddr -> Nonce8 -> STM (Maybe (OnionDestination RouteId))) -- ^ lookup sender of onion query
                -> (UDP.NodeInfo -> RouteId -> IO (Maybe OnionRoute))            -- ^ lookup OnionRoute by id
                -> IO ( ( TVar (ChaChaDRG, Data.Word64Map.Word64Map a)
                        , RelayCache
@@ -375,8 +375,9 @@ newClient :: TransportCrypto
                        , Transport String (OnionDestination RouteId) (OnionMessage Encrypted) )
                      , RelayClient )
 newClient crypto store load lookupSender getRoute = do
+    let runio io = return () -- TODO: run IO action
     (tcpcache,net0) <- toxTCP crypto
-    (relaynet,net1) <- partitionRelay net0
+    (relaynet,net1) <- partitionRelay runio net0
     (onionnet,net2) <- partitionOnion crypto lookupSender getRoute net1
     let net3 = {- XXX: Client type forces this pointless layering. -}
                layerTransport ((Right .) . (,) . (,) False . snd) (,) net2
@@ -428,12 +429,13 @@ showViaRelay (ViaRelay mcon nid tcp) =
     "TCP:" ++ maybe "(oob)" (\(ConId con) -> "(" ++ show con ++ ")") mcon
            ++ show nid ++ "@@" ++ show (nodeAddr tcp)
 
-partitionRelay :: TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket)
+partitionRelay :: (IO () -> STM ())
+                      -> TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket)
                       -> IO ( Transport err ViaRelay ByteString
                             , TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket))
-partitionRelay tr = partitionTransportM parse encode tr
+partitionRelay runio tr = partitionTransportM parse encode tr
   where
-    parse :: ((SessionData,RelayPacket), NodeInfo) -> IO (Either (ByteString, ViaRelay) ((SessionData,RelayPacket),NodeInfo))
+    parse :: ((SessionData,RelayPacket), NodeInfo) -> STM (Either (ByteString, ViaRelay) ((SessionData,RelayPacket),NodeInfo))
     parse ((st,RelayData bs conid), ni) = do
         nid <- getRelayedRemote st conid
         return $ Left (bs, ViaRelay (Just conid) nid ni)
@@ -463,7 +465,7 @@ partitionRelay tr = partitionTransportM parse encode tr
 
 
 partitionOnion :: TransportCrypto
-                  -> (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination RouteId)))
+                  -> (SockAddr -> Nonce8 -> STM (Maybe (OnionDestination RouteId)))
                   -> (UDP.NodeInfo -> RouteId -> IO (Maybe OnionRoute))
                   -> TransportA err NodeInfo (SessionData,RelayPacket) (Bool,RelayPacket)
                   -> IO ( Transport err (OnionDestination RouteId) (OnionMessage Encrypted)
@@ -471,8 +473,8 @@ partitionOnion :: TransportCrypto
 partitionOnion crypto lookupSender getRoute tr = partitionTransportM parse encode tr
  where
     parse :: ((SessionData,RelayPacket), NodeInfo)
-                -> IO (Either (OnionMessage Encrypted   , OnionDestination RouteId)
-                              ((SessionData,RelayPacket), NodeInfo))
+                -> STM (Either (OnionMessage Encrypted   , OnionDestination RouteId)
+                               ((SessionData,RelayPacket), NodeInfo))
     parse pass@((_,OnionPacketResponse msg@(OnionAnnounceResponse n8 _ _)), nodeA) = do
         m <- lookupSender (nodeAddr nodeA) n8
         case m of