Ported dhtd to reworked QueryResponse design.

10 files changed, 147 insertions, 131 deletions

diff --git a/dht/src/Data/Tox/Onion.hs b/dht/src/Data/Tox/Onion.hs
index faff3cdf..1cf89bae 100644
--- a/dht/src/Data/Tox/Onion.hs
+++ b/dht/src/Data/Tox/Onion.hs
@@ -66,8 +66,6 @@ import Util (sameAddress)
 import Text.XXD
 import qualified Data.ByteArray as BA
 
-type HandleLo a = Arrival String SockAddr ByteString -> IO a
-
 type UDPTransport = Transport String SockAddr ByteString
 
 
@@ -186,10 +184,10 @@ onion bs saddr getf = do (f,(asymm,ret3)) <- runGet ((,) <$> getf <*> getOnionRe
                          oaddr <- onionToOwner asymm ret3 saddr
                          return (f asymm, oaddr)
 
-parseOnionAddr :: (SockAddr -> Nonce8 -> IO (Maybe (OnionDestination r)))
+parseOnionAddr :: (SockAddr -> Nonce8 -> STM (Maybe (OnionDestination r)))
                   -> (ByteString, SockAddr)
-                  -> IO (Either (OnionMessage Encrypted,OnionDestination r)
-                                (ByteString,SockAddr))
+                  -> STM (Either (OnionMessage Encrypted,OnionDestination r)
+                                 (ByteString,SockAddr))
 parseOnionAddr lookupSender (msg,saddr)
     | Just (typ,bs) <- B.uncons msg
     , let right = Right (msg,saddr)
@@ -271,24 +269,25 @@ forwardOnions crypto baddr udp sendTCP = udp { awaitMessage = forwardAwait crypt
 
 forwardAwait :: TransportCrypto
                 -> UDPTransport
-                -> (Int -> OnionMessage Encrypted -> IO ()) {- ^ TCP relay send -} -> HandleLo a -> STM (IO a)
-forwardAwait crypto udp sendTCP kont = do
-    fix $ \another0 -> do
-      let another = join $ atomically another0
-      awaitMessage udp $ \case
-        m@(Arrival saddr bs) -> case B.head bs of
-            0x80 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N0) crypto (Addressed saddr) udp another
-            0x81 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N1) crypto (Addressed saddr) udp another
-            0x82 -> forward kont bs $ handleOnionRequest (Proxy :: Proxy  N2) crypto (Addressed saddr) udp another
-            0x8c -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp sendTCP another
-            0x8d -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp sendTCP another
-            0x8e -> forward kont bs $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp sendTCP another
-            _    -> kont m
-        m -> kont m
-
-forward :: (Serialize b, Show b) =>
-              HandleLo a -> ByteString -> (b -> IO a) -> IO a
-forward kont bs f = either (kont . ParseError) f $ decode $ B.tail bs
+                -> (Int -> OnionMessage Encrypted -> IO ()) {- ^ TCP relay send -} -> STM (Arrival String SockAddr ByteString,IO())
+forwardAwait crypto udp sendTCP = do
+    (m,io) <- awaitMessage udp
+    let pass = return (m, io)
+    case m of
+        Arrival saddr bs ->
+          let forward :: Serialize b => (b -> STM (Arrival String SockAddr ByteString, IO ()))
+                                        -> STM (Arrival String SockAddr ByteString, IO ())
+              forward f = either (\e -> return (ParseError e,io)) (fmap (second (io >>)) . f) $ decode $ B.tail bs
+          in case B.head bs of
+            0x80 -> forward $ handleOnionRequest (Proxy :: Proxy  N0) crypto (Addressed saddr) udp
+            0x81 -> forward $ handleOnionRequest (Proxy :: Proxy  N1) crypto (Addressed saddr) udp
+            0x82 -> forward $ handleOnionRequest (Proxy :: Proxy  N2) crypto (Addressed saddr) udp
+            0x8c -> forward $ handleOnionResponse (Proxy :: Proxy N3) crypto saddr udp sendTCP
+            0x8d -> forward $ handleOnionResponse (Proxy :: Proxy N2) crypto saddr udp sendTCP
+            0x8e -> forward $ handleOnionResponse (Proxy :: Proxy N1) crypto saddr udp sendTCP
+            _    -> pass
+        _ -> pass
+
 
 class SumToThree a b
 
@@ -586,28 +585,29 @@ handleOnionRequest :: forall a proxy n.
                       , Sized (ReturnPath n)
                       , Typeable n
                       , Typeable (ThreeMinus (S n))
-                      ) => proxy n -> TransportCrypto -> (forall x. x -> Addressed x) -> UDPTransport -> IO a -> OnionRequest n -> IO a
-handleOnionRequest proxy crypto saddr udp kont (OnionRequest nonce msg rpath) = do
+                      ) => proxy n -> TransportCrypto
+                                   -> (forall x. x -> Addressed x)
+                                   -> UDPTransport
+                                   -> OnionRequest n
+                                   -> STM (Arrival String SockAddr ByteString, IO ())
+handleOnionRequest proxy crypto saddr udp (OnionRequest nonce msg rpath) = do
     let n = peanoVal rpath
-    dput XOnion $ "handleOnionRequest " ++ show n
-    (sym, snonce) <- atomically ( (,) <$> transportSymmetric crypto
-                                      <*> transportNewNonce crypto )
+        io1 = dput XOnion $ "handleOnionRequest " ++ show n
+    (sym, snonce) <- ( (,) <$> transportSymmetric crypto
+                           <*> transportNewNonce crypto )
     peeled <- peelOnion crypto nonce msg
     let showDestination = case saddr () of
             Addressed a _ -> either show show $ either4or6 a
             TCPIndex  i _ -> "TCP" ++ show [i]
 
-    case peeled of
-        Left e -> do
+    fmap (second (io1 >>)) $ case peeled of
+        Left e -> return $ (ParseError e,) $ do
             dput XOnion $ unwords [ "peelOnion:", show n, showDestination, e]
-            kont
-        Right (Addressed dst msg') -> do
+        Right (Addressed dst msg') -> return $ (Discarded,) $ do
             dput XOnion $ unwords [ "peelOnion:", show n, showDestination, "-->", either show show (either4or6 dst), "SUCCESS"]
             sendMessage udp dst (runPut $ putRequest $ OnionRequest nonce msg' $ wrapSymmetric sym snonce saddr rpath)
-            kont
-        Right (TCPIndex {}) -> do
+        Right (TCPIndex {}) -> return $ (,) (ParseError "handleOnionRequest: Onion forward to TCP client?") $ do
             dput XUnexpected "handleOnionRequest: Onion forward to TCP client?"
-            kont
 
 wrapSymmetric :: Serialize (ReturnPath n) =>
                  SymmetricKey -> Nonce24 -> (forall x. x -> Addressed x) -> ReturnPath n -> ReturnPath (S n)
@@ -622,7 +622,7 @@ peelOnion :: ( Typeable n, Typeable t, Serialize (Addressed (Forwarding n t)))
               => TransportCrypto
               -> Nonce24
               -> Forwarding (S n) t
-              -> IO (Either String (Addressed (Forwarding n t)))
+              -> STM (Either String (Addressed (Forwarding n t)))
 peelOnion crypto nonce (Forwarding k fwd) = do
     fmap runIdentity . uncomposed <$> decryptMessage crypto (dhtKey crypto) nonce (Right $ Asymm k nonce fwd)
 
@@ -632,25 +632,22 @@ handleOnionResponse :: (KnownPeanoNat n, Sized (ReturnPath n), Serialize (Return
                         -> SockAddr
                         -> UDPTransport
                         -> (Int -> OnionMessage Encrypted -> IO ()) -- ^ TCP-relay onion send.
-                        -> IO a
                         -> OnionResponse (S n)
-                        -> IO a
-handleOnionResponse proxy crypto saddr udp sendTCP kont (OnionResponse path msg) = do
-    sym <- atomically $ transportSymmetric crypto
+                        -> STM (Arrival String SockAddr ByteString, IO ())
+handleOnionResponse proxy crypto saddr udp sendTCP (OnionResponse path msg) = do
+    sym <- transportSymmetric crypto
     case peelSymmetric sym path of
-        Left e -> do
+        Left e -> return $ (ParseError e,) $ do
             -- todo report encryption error
             let n = peanoVal path
             dput XMisc $ unwords [ "peelSymmetric:", show n, either show show (either4or6 saddr), e]
-            kont
-        Right (Addressed dst path') -> do
+        Right (Addressed dst path') -> return $ (Discarded,) $ do
             sendMessage udp dst (runPut $ putResponse $ OnionResponse path' msg)
-            kont
         Right (TCPIndex dst path') -> do
             case peanoVal path' of
-                0 -> sendTCP dst msg
-                n -> dput XUnexpected $ "handleOnionResponse: TCP-bound OnionResponse" ++ show n ++ " not supported."
-            kont
+                0 -> return (Discarded, sendTCP dst msg)
+                n -> let e = "handleOnionResponse: TCP-bound OnionResponse" ++ show n ++ " not supported."
+                     in return (ParseError e, dput XUnexpected e)
 
 
 data AnnounceRequest = AnnounceRequest
@@ -787,7 +784,7 @@ instance Serialize OnionData where
     put (OnionDHTPublicKey dpk) = put (0x9c :: Word8) >> put dpk
     put (OnionFriendRequest fr) = put (0x20 :: Word8) >> put fr
 
-selectKey :: TransportCrypto -> OnionMessage f -> OnionDestination r -> IO (SecretKey, PublicKey)
+selectKey :: Monad m => TransportCrypto -> OnionMessage f -> OnionDestination r -> m (SecretKey, PublicKey)
 selectKey crypto _ rpath@(OnionDestination (AnnouncingAlias skey pkey) _ _)
                            = return (skey, pkey)
 selectKey crypto msg rpath = return $ aliasKey crypto rpath
@@ -808,32 +805,36 @@ encrypt crypto msg rpath = do
     m <- sequenceMessage $ transcode encipher msg
     return (m, rpath)
 
-decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination r -> IO (Either String (OnionMessage Identity, OnionDestination r))
+decrypt :: TransportCrypto
+            -> OnionMessage Encrypted
+            -> OnionDestination r
+            -> STM (Either String (OnionMessage Identity, OnionDestination r))
 decrypt crypto msg addr = do
     (skey,pkey) <- selectKey crypto msg addr
     let decipher1 :: Serialize a =>
                      TransportCrypto -> SecretKey -> Nonce24
                      -> Either (PublicKey,Encrypted a) (Asymm (Encrypted a))
-                     -> (IO ∘ Either String ∘ Identity) a
+                     -> (STM ∘ Either String ∘ Identity) a
         decipher1 crypto k n arg = Composed $ do
             let (sender,e) = either id (senderKey &&& asymmData) arg
-            secret <- lookupSharedSecret crypto k sender n
+            secret <- lookupSharedSecretSTM crypto k sender n
             return $ Composed $ do
                 plain <- ToxCrypto.decrypt secret e
                 Identity <$> decodePlain plain
         decipher :: Serialize a
             => Nonce24 -> Either (Encrypted a) (Asymm (Encrypted a))
-                       -> (IO ∘ Either String ∘ Identity) a
+                       -> (STM ∘ Either String ∘ Identity) a
         decipher = (\n -> decipher1 crypto skey n . left (senderkey addr))
     foo <- sequenceMessage $ transcode decipher msg
     let result = do
             msg <- sequenceMessage foo
             Right (msg, addr)
-    case msg of
-        OnionToRouteResponse {} -> case result of
-            Left e  -> dput XMan $ "Error decrypting data-to-route response: " ++ e
-            Right m -> dput XMan $ "Decrypted data-to-route response: " ++ show (fst m)
-        _ -> return ()
+    -- -- TODO runio
+    -- case msg of
+    --     OnionToRouteResponse {} -> case result of
+    --         Left e  -> dput XMan $ "Error decrypting data-to-route response: " ++ e
+    --         Right m -> dput XMan $ "Decrypted data-to-route response: " ++ show (fst m)
+    --     _ -> return ()
     return result
 
 senderkey :: OnionDestination r -> t -> (PublicKey, t)
@@ -857,11 +858,11 @@ decryptMessage :: (Typeable x, Serialize x) =>
                   -> Nonce24
                   -> Either (PublicKey, Encrypted x)
                             (Asymm (Encrypted x))
-                  -> IO ((Either String ∘ Identity) x)
+                  -> STM ((Either String ∘ Identity) x)
 decryptMessage crypto (sk,pk) n arg = do
     let (sender,e) = either id (senderKey &&& asymmData) arg
         plain      = Composed . fmap Identity . (>>= decodePlainVerbose)
-    secret <- lookupSharedSecret crypto sk sender n
+    secret <- lookupSharedSecretSTM crypto sk sender n
     return $ plain $ ToxCrypto.decrypt secret e
 
 sequenceMessage :: Applicative m => OnionMessage (m ∘ f) -> m (OnionMessage f)
@@ -1002,9 +1003,9 @@ selectAlias crypto pkey = do
 parseDataToRoute
     :: TransportCrypto
          -> (OnionMessage Encrypted,OnionDestination r)
-         -> IO (Either ((PublicKey,OnionData),AnnouncedRendezvous) (OnionMessage Encrypted, OnionDestination r))
+         -> STM (Either ((PublicKey,OnionData),AnnouncedRendezvous) (OnionMessage Encrypted, OnionDestination r))
 parseDataToRoute crypto (OnionToRouteResponse dta, od) = do
-  ks <- atomically $ userKeys crypto
+  ks <- userKeys crypto
 
   omsg0 <- decryptMessage crypto (rendezvousSecret crypto,rendezvousPublic crypto)
                                 (asymmNonce dta)
@@ -1035,17 +1036,19 @@ parseDataToRoute crypto (OnionToRouteResponse dta, od) = do
                     (dataFromKey dtr)
                     $ Rendezvous (rendezvousPublic crypto) $ onionNodeInfo od )
       r = either (const $ Right (OnionToRouteResponse dta,od)) Left e
-  -- parseDataToRoute OnionToRouteResponse decipherAndAuth: auth fail
-  case e of
-    Left _  -> dput XMisc $ "Failed keys: " ++ show (map (key2id . snd) ks)
-    Right _ -> return ()
-  dput XMisc $ unlines
-        [ "parseDataToRoute " ++ either id (const "Right") e
-        , "       crypto inner.me   = " ++ either id (\(pk,_,_) -> show $ key2id pk) eInner
-        , "              inner.them = " ++ either id (show . key2id . dataFromKey) eOuter
-        , "              outer.me   = " ++ show (key2id $ rendezvousPublic crypto)
-        , "              outer.them = " ++ show (key2id $ senderKey dta)
-        ]
+      io :: IO ()
+      io = do
+          case e of
+            Left _  -> dput XMisc $ "Failed keys: " ++ show (map (key2id . snd) ks)
+            Right _ -> return ()
+          dput XMisc $ unlines
+                [ "parseDataToRoute " ++ either id (const "Right") e
+                , "       crypto inner.me   = " ++ either id (\(pk,_,_) -> show $ key2id pk) eInner
+                , "              inner.them = " ++ either id (show . key2id . dataFromKey) eOuter
+                , "              outer.me   = " ++ show (key2id $ rendezvousPublic crypto)
+                , "              outer.them = " ++ show (key2id $ senderKey dta)
+                ]
+  -- TODO: run io
   return r
 parseDataToRoute _ msg = return $ Right msg
 
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