1 files changed, 7 insertions, 5 deletions

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 ())