9 files changed, 72 insertions, 59 deletions

diff --git a/dht/src/Network/BitTorrent/MainlineDHT.hs b/dht/src/Network/BitTorrent/MainlineDHT.hs
index 705d7291..e0715d4a 100644
--- a/dht/src/Network/BitTorrent/MainlineDHT.hs
+++ b/dht/src/Network/BitTorrent/MainlineDHT.hs
@@ -431,11 +431,11 @@ showPacket f addr flow bs = L8.unpack $ L8.unlines es
 -- Add detailed printouts for every packet.
 addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
 addVerbosity tr =
-    tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
+    tr { awaitMessage = do
+            (m,io) <- awaitMessage tr
             case m of
-                Arrival addr msg -> dput XBitTorrent (showPacket id addr " --> " msg)
-                _                -> return ()
-            kont m
+                Arrival addr msg -> return (m, io >> dput XBitTorrent (showPacket id addr " --> " msg))
+                _                -> return (m, io)
        , sendMessage = \addr msg -> do
             dput XBitTorrent (showPacket id addr " <-- " msg)
             sendMessage tr addr msg
@@ -603,7 +603,7 @@ newClient swarms addr udp = do
         -- recursive since 'updateRouting' does not invoke 'awaitMessage' which
         -- which was modified by 'onInbound'.  However, I'm going to avoid the
         -- mutual reference just to be safe.
-        outgoingClient = client { clientNet = net { awaitMessage = pure . ($ Terminated) } }
+        outgoingClient = client { clientNet = net { awaitMessage = pure (Terminated, return ()) } }
 
         dispatch = DispatchMethods
             { classifyInbound = classify -- :: x -> MessageClass err meth tid addr x
diff --git a/dht/src/Network/Lossless.hs b/dht/src/Network/Lossless.hs
index 41203ca5..7ccceec1 100644
--- a/dht/src/Network/Lossless.hs
+++ b/dht/src/Network/Lossless.hs
@@ -61,7 +61,9 @@ lossless lbl isLossless encode saddr udp = do
     rloop <- forkIO $ do
         -- This thread enqueues inbound packets or writes them to the oob
         -- channel.
-        fix $ \loop -> join $ atomically $ awaitMessage udp $ \m -> do
+        fix $ \loop -> do
+            (m,io) <- atomically $ awaitMessage udp
+            io
             m' <- case m of Terminated   -> return Nothing
                             ParseError e -> return $ Just (Left e)
                             Arrival a x  -> Just . Right <$> isLossless x a
@@ -87,15 +89,14 @@ lossless lbl isLossless encode saddr udp = do
                             -- we will use this STM action stop it from waiting on the oob TChan.
                             -- XXX: This shouldn't be neccessary and might be costly.
     let tr = Transport
-            { awaitMessage = \kont ->
+            { awaitMessage =
                 orElse
                     (do x <- readTChan oob `orElse` join (readTVar term)
-                        return $ kont $! x)
+                        return (x, return ()))
                     (do x <- PB.awaitReadyPacket pb
                         report <- pbReport "dequeued" pb
-                        return $ do
-                            atomically $ writeTChan oob (ParseError report)
-                            kont $! uncurry (flip Arrival) x)
+                        return $ (,) (uncurry (flip Arrival) x) $ do
+                            atomically $ writeTChan oob (ParseError report))
             , sendMessage = \a' x' -> do
                 seqno <- atomically $ do
                     seqno <- PB.nextToSendSequenceNumber pb
diff --git a/dht/src/Network/SessionTransports.hs b/dht/src/Network/SessionTransports.hs
index b6d02f36..68233cd4 100644
--- a/dht/src/Network/SessionTransports.hs
+++ b/dht/src/Network/SessionTransports.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE LambdaCase     #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE TupleSections  #-}
 module Network.SessionTransports
     ( Sessions
     , initSessions
@@ -73,9 +74,9 @@ newSession Sessions{sessionsByAddr,sessionsById,sessionIds,sessionsSendRaw} unwr
             return sid
     forM msid $ \sid -> do
       let tr = Transport
-            { awaitMessage = \kont -> do
+            { awaitMessage = do
                 x <- takeTMVar mvar
-                return $ kont $! maybe Terminated (uncurry $ flip Arrival) x
+                return $ (, return ()) $ maybe Terminated (uncurry $ flip Arrival) x
             , sendMessage = \addr x -> do
                 x' <- unwrap addr x
                 sessionsSendRaw saddr x'
@@ -92,8 +93,9 @@ newSession Sessions{sessionsByAddr,sessionsById,sessionIds,sessionsSendRaw} unwr
             }
       return (sid,tr)
 
-sessionHandler :: Sessions x -> (Multi.SessionAddress -> x -> IO (Maybe (x -> x)))
-sessionHandler Sessions{sessionsByAddr} = \addr0 x -> do
+sessionHandler :: Sessions x -> Arrival err Multi.SessionAddress x
+                             -> STM (Arrival err Multi.SessionAddress x, IO ())
+sessionHandler Sessions{sessionsByAddr} (Arrival addr0 x) = return $ (,) Discarded $ do
     let addr = -- Canonical in case of 6-mapped-4 addresses.
                Multi.canonize addr0
         dispatch []     = return ()
@@ -101,4 +103,4 @@ sessionHandler Sessions{sessionsByAddr} = \addr0 x -> do
                              when (not b) $ dispatch fs
     fs <- atomically $ Map.lookup addr <$> readTVar sessionsByAddr
     mapM_ (dispatch . IntMap.elems) fs
-    return Nothing -- consume all packets.
+sessionHandler _ m = return (m, return ())
diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index 6b39d57a..3dd1d48e 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -245,14 +245,15 @@ isLocalHost _                       = False
 
 addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
 addVerbosity tr =
-    tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
+    tr { awaitMessage = do
+            (m,io) <- awaitMessage tr
             case m of
-              Arrival addr msg -> do
+              Arrival addr msg -> return $ (,) m $ do
+                io
                 when (not (B.null msg || elem (B.head msg) [0,1,2,4,0x81,0x82,0x8c,0x8d])) $ do
                     mapM_ (\x -> dput XMisc ( (show addr) ++ " --> " ++ x))
                           $ xxd 0 msg
-              _ -> return ()
-            kont m
+              _ -> return (m,io)
        , sendMessage = \addr msg -> do
             when (not (B.null msg || elem (B.head msg) [0,1,2,4,0x81,0x8c,0x8d])) $ do
                 mapM_ (\x -> dput XMisc ( (show addr) ++ " <-- " ++ x))
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