STM-based awaitMessage.

9 files changed, 112 insertions, 81 deletions

diff --git a/dht/src/Data/Tox/Onion.hs b/dht/src/Data/Tox/Onion.hs
index d3c8086d..e0d7c744 100644
--- a/dht/src/Data/Tox/Onion.hs
+++ b/dht/src/Data/Tox/Onion.hs
@@ -60,7 +60,7 @@ import Data.Bits (shiftR,shiftL)
 import qualified Rank2
 import Util (sameAddress)
 
-type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a
+type HandleLo a = Arrival String SockAddr ByteString -> IO a
 
 type UDPTransport = Transport String SockAddr ByteString
 
@@ -264,11 +264,12 @@ forwardOnions crypto baddr udp sendTCP = udp { awaitMessage = forwardAwait crypt
 
 forwardAwait :: TransportCrypto
                 -> UDPTransport
-                -> (Int -> OnionMessage Encrypted -> IO ()) {- ^ TCP relay send -} -> HandleLo a -> IO a
+                -> (Int -> OnionMessage Encrypted -> IO ()) {- ^ TCP relay send -} -> HandleLo a -> STM (IO a)
 forwardAwait crypto udp sendTCP kont = do
-    fix $ \another -> do
+    fix $ \another0 -> do
+      let another = join $ atomically another0
       awaitMessage udp $ \case
-        m@(Just (Right (bs,saddr))) -> case B.head bs of
+        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
@@ -278,9 +279,9 @@ forwardAwait crypto udp sendTCP kont = do
             _    -> kont m
         m -> kont m
 
-forward :: forall c b b1. (Serialize b, Show b) =>
-              (Maybe (Either String b1) -> c) -> ByteString -> (b -> c) -> c
-forward kont bs f = either (kont . Just . Left) f $ decode $ B.tail bs
+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
 
 class SumToThree a b
 
diff --git a/dht/src/Network/BitTorrent/MainlineDHT.hs b/dht/src/Network/BitTorrent/MainlineDHT.hs
index ed97ee31..0269268f 100644
--- a/dht/src/Network/BitTorrent/MainlineDHT.hs
+++ b/dht/src/Network/BitTorrent/MainlineDHT.hs
@@ -428,8 +428,9 @@ showPacket f addr flow bs = L8.unpack $ L8.unlines es
 addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
 addVerbosity tr =
     tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
-            forM_ m $ mapM_ $ \(msg,addr) -> do
-                dput XBitTorrent (showPacket id addr " --> " msg)
+            case m of
+                Arrival addr msg -> dput XBitTorrent (showPacket id addr " --> " msg)
+                _                -> return ()
             kont m
        , sendMessage = \addr msg -> do
             dput XBitTorrent (showPacket id addr " <-- " msg)
@@ -598,7 +599,7 @@ newClient swarms addr = 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 = ($ Nothing) } }
+        outgoingClient = client { clientNet = net { awaitMessage = pure . ($ Terminated) } }
 
         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 861792ab..5a313aed 100644
--- a/dht/src/Network/Lossless.hs
+++ b/dht/src/Network/Lossless.hs
@@ -59,37 +59,38 @@ lossless isLossless encode saddr udp = do
     rloop <- forkIO $ do
         -- This thread enqueues inbound packets or writes them to the oob
         -- channel.
-        myThreadId >>= flip labelThread ("lossless."++show saddr)
-        fix $ \loop -> do
-        awaitMessage udp $ \m -> do
-            m' <- mapM (mapM $ uncurry isLossless) m
+        fix $ \loop -> join $ atomically $ awaitMessage udp $ \m -> do
+            m' <- case m of Terminated   -> return Nothing
+                            ParseError e -> return $ Just (Left e)
+                            Arrival a x  -> Just . Right <$> isLossless x a
             case m' of
                 Nothing -> do
-                    atomically $ writeTChan oob Nothing
+                    atomically $ writeTChan oob Terminated
                     -- Quit thread here.
                 Just (Left e) -> do
-                    atomically $ writeTChan oob (Just $ Left e)
+                    atomically $ writeTChan oob (ParseError e)
                     loop
                 Just (Right event) -> do
                     atomically $ do
                         -- x' <- isLossless xaddr x
                         PB.grokInboundPacket pb event
                         case event of
-                            PacketReceivedLossy {} -> writeTChan oob (Just $ Right $ peReceivedPayload event)
+                            PacketReceivedLossy {} -> writeTChan oob (uncurry (flip Arrival) $ peReceivedPayload event)
                             _                      -> do
                                 report <- pbReport "enqueued" pb
-                                writeTChan oob (Just $ Left report)
+                                writeTChan oob (ParseError report)
                     loop
+    labelThread rloop ("lossless."++show saddr)
     let tr = Transport
-            { awaitMessage = \kont -> do
-                join $ atomically $ orElse
+            { awaitMessage = \kont ->
+                orElse
                     (do x <- readTChan oob
                         return $ kont $! x)
                     (do x <- PB.awaitReadyPacket pb
                         report <- pbReport "dequeued" pb
                         return $ do
-                            atomically $ writeTChan oob (Just $ Left report)
-                            kont $! Just (Right x))
+                            atomically $ writeTChan oob (ParseError report)
+                            kont $! uncurry (flip Arrival) x)
             , sendMessage = \a' x' -> do
                 seqno <- atomically $ do
                     seqno <- PB.nextToSendSequenceNumber pb
@@ -111,9 +112,11 @@ lossless isLossless encode saddr udp = do
                         when isfull retry
                 let sendit = sendMessage udp saddr x
                 maybe sendit (catchIOError sendit) oops
-            , closeTransport = do
-                atomically $ writeTChan oob Nothing -- quit rloop thread
-                closeTransport udp
+            , setActive = \case
+                False -> do
+                    atomically $ writeTChan oob Terminated -- quit rloop thread
+                    setActive udp False
+                True -> return ()
             }
         resend ns = do
             xs <- atomically $ retrieveForResend pb ns
diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index f62fbefe..ba686929 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -54,15 +54,18 @@ data Arrival err addr x
 data TransportA err addr x y = Transport
     { -- | Blocks until an inbound packet is available. Then calls the provided
       -- continuation with the packet and origin addres or an error condition.
-      awaitMessage :: forall a. (Arrival err addr x -> IO a) -> IO a
+      awaitMessage :: forall a. (Arrival err addr x -> IO a) -> STM (IO a)
       -- | Send an /y/ packet to the given destination /addr/.
     , sendMessage :: addr -> y -> IO ()
       -- | Shutdown and clean up any state related to this 'Transport'.
-    , closeTransport :: IO ()
+    , setActive :: Bool -> IO ()
     }
 
 type Transport err addr x = TransportA err addr x x
 
+closeTransport :: TransportA err addr x y -> IO ()
+closeTransport tr = setActive tr False
+
 -- | This function modifies a 'Transport' to use higher-level addresses and
 -- packet representations.  It could be used to change UDP 'ByteString's into
 -- bencoded syntax trees or to add an encryption layer in which addresses have
@@ -110,7 +113,7 @@ layerTransport parse encode tr =
                     (\x' addr' -> return $ encode x' addr')
                     tr
 
--- | Paritions a 'Transport' into two higher-level transports.  Note: An 'MVar'
+-- | Paritions a 'Transport' into two higher-level transports.  Note: A 'TChan'
 -- is used to share the same underlying socket, so be sure to fork a thread for
 -- both returned 'Transport's to avoid hanging.
 partitionTransportM :: ((b,a) -> IO (Either (x,xaddr) (b,a)))
@@ -118,28 +121,28 @@ partitionTransportM :: ((b,a) -> IO (Either (x,xaddr) (b,a)))
                       -> Transport err a b
                       -> IO (Transport err xaddr x, Transport err a b)
 partitionTransportM parse encodex tr = do
-    mvar <- newEmptyMVar
+    tchan <- atomically newTChan
     let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
                     awaitMessage tr $ \m -> case m of
                         Arrival adr msg  -> parse (msg,adr) >>= \case
                             Left (x,xaddr) -> kont $ Arrival xaddr x
-                            Right y        -> putMVar mvar (Just y) >> again
+                            Right y        -> atomically (writeTChan tchan (Just y)) >> join (atomically again)
                         ParseError e     -> kont $ ParseError e
-                        Terminated       -> putMVar mvar Nothing >> kont Terminated
+                        Terminated       -> atomically (writeTChan tchan Nothing) >> kont Terminated
                  , sendMessage = \addr' msg' -> do
                     msg_addr <- encodex (msg',addr')
                     mapM_ (uncurry . flip $ sendMessage tr) msg_addr
                  }
         ytr = Transport
-                { awaitMessage = \kont -> takeMVar mvar >>= kont . \case
+                { awaitMessage = \kont -> readTChan tchan >>= pure . kont . \case
                     Nothing        -> Terminated
                     Just (y,yaddr) -> Arrival yaddr y
                 , sendMessage = sendMessage tr
-                , closeTransport = return ()
+                , setActive = const $ return ()
                 }
     return (xtr, ytr)
 
--- | Paritions a 'Transport' into two higher-level transports.  Note: An 'MVar'
+-- | Paritions a 'Transport' into two higher-level transports.  Note: An 'TChan'
 -- is used to share the same underlying socket, so be sure to fork a thread for
 -- both returned 'Transport's to avoid hanging.
 partitionTransport :: ((b,a) -> Either (x,xaddr) (b,a))
@@ -157,14 +160,14 @@ partitionAndForkTransport ::
                       -> Transport err a b
                       -> IO (Transport err xaddr x, Transport err a b)
 partitionAndForkTransport forkedSend parse encodex tr = do
-    mvar <- newEmptyMVar
+    tchan <- atomically newTChan
     let xtr = tr { awaitMessage = \kont -> fix $ \again -> do
                     awaitMessage tr $ \case
                         Arrival a b -> parse (b,a) >>= \case
                             Left (x,xaddr) -> kont $ Arrival xaddr x
-                            Right (b,a)    -> putMVar mvar (Arrival a b) >> again
+                            Right (b,a)    -> atomically (writeTChan tchan (Arrival a b)) >> join (atomically again)
                         ParseError e     -> kont $ ParseError e
-                        Terminated       -> putMVar mvar Terminated >> kont Terminated
+                        Terminated       -> atomically (writeTChan tchan Terminated) >> kont Terminated
                  , sendMessage = \addr' msg' -> do
                     msg_addr <- encodex (msg',addr')
                     case msg_addr of
@@ -173,9 +176,9 @@ partitionAndForkTransport forkedSend parse encodex tr = do
                         Nothing               -> return ()
                  }
         ytr = Transport
-                { awaitMessage   = \kont -> takeMVar mvar >>= kont
+                { awaitMessage   = \kont -> readTChan tchan >>= pure . kont
                 , sendMessage    = sendMessage tr
-                , closeTransport = return ()
+                , setActive      = \_ -> return ()
                 }
     return (xtr, ytr)
 
@@ -186,7 +189,7 @@ partitionAndForkTransport forkedSend parse encodex tr = do
 addHandler :: (err -> IO ()) -> (addr -> x -> IO (Maybe (x -> x))) -> Transport err addr x -> Transport err addr x
 addHandler onParseError f tr = tr
     { awaitMessage = \kont -> fix $ \eat -> awaitMessage tr $ \case
-        Arrival addr x -> f addr x >>= maybe eat (kont . Arrival addr . ($ x))
+        Arrival addr x -> f addr x >>= maybe (join $ atomically eat) (kont . Arrival addr . ($ x))
         ParseError e   -> onParseError e >> kont (ParseError e)
         Terminated     -> kont Terminated
     }
@@ -210,14 +213,15 @@ onInbound f tr = addHandler (const $ return ()) (\addr x -> f addr x >> return (
 --  > quitServer
 forkListener :: String -> Transport err addr x -> IO (IO ())
 forkListener name client = do
+    setActive client True
     thread_id <- forkIO $ do
         myThreadId >>= flip labelThread ("listener."++name)
-        fix $ \loop -> awaitMessage client $ \case
+        fix $ \loop -> join $ atomically $ awaitMessage client $ \case
             Terminated -> return ()
             _          -> loop
         dput XMisc $ "Listener died: " ++ name
     return $ do
-        closeTransport client
+        setActive client False
         -- killThread thread_id
 
 -- * Implementing a query\/response 'Client'.
@@ -606,11 +610,11 @@ udpTransport' bind_address = do
   setSocketOption sock Broadcast 1
   bind sock bind_address
   isClosed <- newEmptyMVar
+  udpTChan <- atomically newTChan
   let tr = Transport {
       awaitMessage = \kont -> do
-        r <- handle (ignoreEOF sock isClosed $ Arrival (SockAddrInet 0 0) B.empty) $ do
-                uncurry (flip Arrival) <$!> B.recvFrom sock udpBufferSize
-        kont $! r
+        r <- readTChan udpTChan
+        return $ kont $! r
     , sendMessage = case family of
         AF_INET6 -> \case
             (SockAddrInet port addr) -> \bs ->
@@ -626,7 +630,8 @@ udpTransport' bind_address = do
             addr@(SockAddrInet6 {}) -> \bs -> dput XMisc ("Discarding packet to "++show addr)
             addr4 -> \bs -> saferSendTo sock bs addr4
         _ -> \addr bs -> saferSendTo sock bs addr
-    , closeTransport = do
+    , setActive = \case
+      False -> do
         dput XMisc $ "closeTransport for udpTransport' called. " ++ show bind_address
         tryPutMVar isClosed () -- signal awaitMessage that the transport is closed.
 #if MIN_VERSION_network (3,1,0)
@@ -639,6 +644,14 @@ udpTransport' bind_address = do
             let sorryGHCButIAmNotFuckingClosingTheSocketYet fd = return ()
             -- This call is necessary to interrupt the blocking recvFrom call in awaitMessage.
             closeFdWith sorryGHCButIAmNotFuckingClosingTheSocketYet (fromIntegral fd)
+      True -> do
+            udpThread <- forkIO $ fix $ \again -> do
+                r <- handle (ignoreEOF sock isClosed $ Arrival (SockAddrInet 0 0) B.empty) $ do
+                        uncurry (flip Arrival) <$!> B.recvFrom sock udpBufferSize
+                atomically $ writeTChan udpTChan r
+                case r of Terminated -> return ()
+                          _          -> again
+            labelThread udpThread ("udp.io."++show bind_address)
     }
   return (tr, sock)
 
@@ -652,13 +665,15 @@ udpTransport bind_address = fst <$> udpTransport' bind_address
 chanTransport :: (addr -> TChan (x, addr)) -> addr -> TChan (x, addr) -> TVar Bool -> Transport err addr x
 chanTransport chanFromAddr self achan aclosed = Transport
     { awaitMessage = \kont -> do
-        x <- atomically $ (uncurry (flip Arrival) <$> readTChan achan)
+        x <- (uncurry (flip Arrival) <$> readTChan achan)
                             `orElse`
-                          (readTVar aclosed >>= check >> return Terminated)
-        kont x
+             (readTVar aclosed >>= check >> return Terminated)
+        return $ kont x
     , sendMessage = \them bs -> do
         atomically $ writeTChan (chanFromAddr them) (bs,self)
-    , closeTransport = atomically $ writeTVar aclosed True
+    , setActive = \case
+        False -> atomically $ writeTVar aclosed True
+        True  -> return ()
     }
 
 -- | Returns a pair of transports linked together to simulate two computers talking to each other.
diff --git a/dht/src/Network/QueryResponse/TCP.hs b/dht/src/Network/QueryResponse/TCP.hs
index 67c19512..0028a5b6 100644
--- a/dht/src/Network/QueryResponse/TCP.hs
+++ b/dht/src/Network/QueryResponse/TCP.hs
@@ -1,5 +1,6 @@
+{-# LANGUAGE CPP                        #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE CPP #-}
+{-# LANGUAGE LambdaCase                 #-}
 module Network.QueryResponse.TCP where
 
 #ifdef THREAD_DEBUG
@@ -11,6 +12,7 @@ import GHC.Conc                  (labelThread)
 
 import Control.Arrow
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TMVar
 import Control.Monad
 import Data.ByteString (ByteString,hPut)
 import Data.Function
@@ -78,7 +80,7 @@ showStat r = case r of PendingTCPSession -> "pending."
 tcp_timeout :: Int
 tcp_timeout = 10000000
 
-acquireConnection :: MVar (Maybe (Either a (x, addr)))
+acquireConnection :: TMVar (Arrival a addr x)
                            -> TCPCache (SessionProtocol x y)
                            -> StreamHandshake addr x y
                            -> addr
@@ -131,10 +133,10 @@ acquireConnection mvar tcpcache stream addr bDoCon = do
                     dput XTCP $ "TCP streamDecode " ++ show (streamAddr stream addr) ++ " --> " ++ maybe "Nothing" (const "got") x
                     case x of
                         Just u -> do
-                            m <- timeout tcp_timeout $ putMVar mvar $ Just $ Right (u, addr)
+                            m <- timeout tcp_timeout $ atomically (putTMVar mvar $ Arrival addr u)
                             when (isNothing m) $ do
                                 dput XTCP $ "TCP "++show (streamAddr stream addr) ++ " dropped packet."
-                                tryTakeMVar mvar
+                                atomically $ tryTakeTMVar mvar
                                 return ()
                             loop
                         Nothing -> do
@@ -206,14 +208,16 @@ tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
                 -> StreamHandshake addr x y
                 -> IO (TCPCache (SessionProtocol x y), TransportA err addr x (Bool,y))
 tcpTransport maxcon stream = do
-    msgvar <- newEmptyMVar
+    msgvar <- atomically newEmptyTMVar
     tcpcache <- atomically $ (`TCPCache` maxcon) <$> newTVar (MM.empty)
     return $ (,) tcpcache Transport
-        { awaitMessage   = \f -> takeMVar msgvar >>= \x -> f x `catchIOError` (\e -> dput XTCP ("TCP transport stopped. " ++ show e) >> f Nothing)
+        { awaitMessage   = \f -> takeTMVar msgvar >>= \x -> return $ do
+            f x `catchIOError` (\e -> dput XTCP ("TCP transport stopped. " ++ show e) >> f Terminated)
         , sendMessage    = \addr (bDoCon,y) -> do
                                 void . forkLabeled "tcp-send" $ do
                                     msock <- acquireConnection msgvar tcpcache stream addr bDoCon
                                     mapM_ ($ y) msock
                                   `catchIOError` \e -> dput XTCP $ "TCP-send: " ++ show e
-        , closeTransport = closeAll tcpcache stream >> putMVar msgvar Nothing
+        , setActive      = \case False -> closeAll tcpcache stream >> atomically (putTMVar msgvar Terminated)
+                                 True  -> return ()
         }
diff --git a/dht/src/Network/SessionTransports.hs b/dht/src/Network/SessionTransports.hs
index e9daf6c1..b36fbcfd 100644
--- a/dht/src/Network/SessionTransports.hs
+++ b/dht/src/Network/SessionTransports.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE LambdaCase     #-}
 {-# LANGUAGE NamedFieldPuns #-}
 module Network.SessionTransports
     ( Sessions
@@ -8,6 +9,7 @@ module Network.SessionTransports
 
 import Control.Concurrent
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TMVar
 import Control.Monad
 import qualified Data.IntMap.Strict as IntMap
          ;import Data.IntMap.Strict (IntMap)
@@ -51,14 +53,14 @@ newSession :: Sessions raw
                     -> SockAddr
                     -> IO (Maybe (Int,TransportA err addr x y))
 newSession Sessions{sessionsByAddr,sessionsById,sessionIds,sessionsSendRaw} unwrap wrap addr0 = do
-    mvar <- newEmptyMVar
+    mvar <- atomically newEmptyTMVar
     let saddr = -- Canonical in case of 6-mapped-4 addresses.
                either id id $ either4or6 addr0
         handlePacket x = do
             m <- wrap saddr x
             case m of
                 Nothing -> return False
-                Just x' -> do putMVar mvar $! Just $! x'
+                Just x' -> do atomically $ putTMVar mvar $! Just $! x'
                               return True
     msid <- atomically $ do
         msid <- S.nearestOutsider 0 <$> readTVar sessionIds
@@ -69,22 +71,25 @@ newSession Sessions{sessionsByAddr,sessionsById,sessionIds,sessionsSendRaw} unwr
                                        $ IntMap.singleton sid handlePacket
             return sid
     forM msid $ \sid -> do
-    let tr = Transport
+      let tr = Transport
             { awaitMessage = \kont -> do
-                x <- takeMVar mvar
-                kont $! Right <$> x
+                x <- takeTMVar mvar
+                return $ kont $! maybe Terminated (uncurry $ flip Arrival) x
             , sendMessage = \addr x -> do
                 x' <- unwrap addr x
                 sessionsSendRaw saddr x'
-            , closeTransport = do
-                tryTakeMVar mvar
-                putMVar mvar Nothing
-                atomically $ do
-                    modifyTVar' sessionIds $ S.delete sid
-                    modifyTVar' sessionsById $ IntMap.delete sid
-                    modifyTVar' sessionsByAddr $ Map.alter (rmSession sid) saddr
+            , setActive = \case
+                False -> do
+                    atomically $ do
+                        tryTakeTMVar mvar
+                        putTMVar mvar Nothing
+                    atomically $ do
+                        modifyTVar' sessionIds $ S.delete sid
+                        modifyTVar' sessionsById $ IntMap.delete sid
+                        modifyTVar' sessionsByAddr $ Map.alter (rmSession sid) saddr
+                True -> return ()
             }
-    return (sid,tr)
+      return (sid,tr)
 
 sessionHandler :: Sessions x -> (SockAddr -> x -> IO (Maybe (x -> x)))
 sessionHandler Sessions{sessionsByAddr} = \addr0 x -> do
diff --git a/dht/src/Network/Tox.hs b/dht/src/Network/Tox.hs
index 61a1d117..4b66cfc8 100644
--- a/dht/src/Network/Tox.hs
+++ b/dht/src/Network/Tox.hs
@@ -222,10 +222,12 @@ isLocalHost _                       = False
 addVerbosity :: Transport err SockAddr ByteString -> Transport err SockAddr ByteString
 addVerbosity tr =
     tr { awaitMessage = \kont -> awaitMessage tr $ \m -> do
-            forM_ m $ mapM_ $ \(msg,addr) -> do
+            case m of
+              Arrival addr msg -> do
                 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
        , sendMessage = \addr msg -> do
             when (not (B.null msg || elem (B.head msg) [0,1,2,4,0x81,0x8c,0x8d])) $ do
diff --git a/dht/src/Network/Tox/AggregateSession.hs b/dht/src/Network/Tox/AggregateSession.hs
index 8c728660..999c7399 100644
--- a/dht/src/Network/Tox/AggregateSession.hs
+++ b/dht/src/Network/Tox/AggregateSession.hs
@@ -188,21 +188,21 @@ forkSession c s setStatus = forkIO $ do
             now <- getPOSIXTime
             atomically $ modifyTVar' q $ PSQ.insert (fromEnum DoTimeout) (now + 15)
 
-        onPacket body loop Nothing          = return ()
-        onPacket body loop (Just (Left e))  = inPrint e >> loop
-        onPacket body loop (Just (Right x)) = body loop x
+        onPacket body loop Terminated     = return ()
+        onPacket body loop (ParseError e) = inPrint e >> loop
+        onPacket body loop (Arrival _ x)  = body loop x
 
-        awaitPacket body = fix $ awaitMessage (sTransport s) . onPacket body
+        awaitPacket body = fix $ join . atomically . awaitMessage (sTransport s) . onPacket body
 
     atomically $ setStatus $ InProgress AwaitingSessionPacket
-    awaitPacket $ \_ (online,()) -> do
+    awaitPacket $ \_ online -> do
         when (msgID online /= M ONLINE) $ do
             inPrint $ "Unexpected initial packet: " ++ show (msgID online)
         atomically $ do setStatus Established
                         sendPacket online
         bump
         beacon <- forkIO $ keepAlive s q
-        awaitPacket $ \awaitNext (x,()) -> do
+        awaitPacket $ \awaitNext x -> do
             bump
             case msgID x of
                 M ALIVE      -> return ()
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)