2 files changed, 31 insertions, 18 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 8fd1402d..6e2647d1 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -990,12 +990,12 @@ noArgPing f [] x = f x
 noArgPing _ _  _ = return Nothing
 
 newXmmpSource :: Tox.NetCryptoSession -> C.Source IO Tox.CryptoMessage
-newXmmpSource sessions = do
-        v <- liftIO $ _todo sessions {- receive a fucking message -}
-        case v of
-            Nothing            -> return () -- Nothing indicates EOF.
-            Just cryptomessage -> do C.yield cryptomessage
-                                     newXmmpSource sessions
+newXmmpSource session = do
+    v <- liftIO $ Tox.receiveCrypto session
+    case v of
+        Nothing            -> return () -- Nothing indicates EOF.
+        Just cryptomessage -> do C.yield cryptomessage
+                                 newXmmpSource session
 
 newXmmpSink :: Tox.NetCryptoSession -> C.Sink (Flush Tox.CryptoMessage) IO ()
 newXmmpSink sessions = C.awaitForever $ \flush_cyptomessage -> do
diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index d43535ad..e6669c3e 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -29,6 +29,7 @@ import qualified Data.Set         as Set
 import qualified Data.Array.Unboxed as A
 import SensibleDir
 import System.FilePath
+import System.IO
 import System.IO.Temp
 import System.Environment
 import System.Directory
@@ -102,7 +103,6 @@ data NetCryptoSession = NCrypto { ncState          :: TVar NetCryptoSessionStatu
                                 , ncView :: TVar SessionView
                                 , ncPacketQueue :: PacketQueue CryptoData
                                 , ncBufferStart :: TVar Word32
-                                , ncDequeueThread :: Maybe ThreadId
                                 , ncPingMachine :: Maybe PingMachine
                                 , ncOutgoingQueue :: PQ.PacketOutQueue (State,Nonce24,TVar MsgOutMap) CryptoMessage (CryptoPacket Encrypted) CryptoData
                                 }
@@ -266,6 +266,19 @@ ncToWire getState seqno bufend pktno msg = do
                                         pkt = CryptoPacket { pktNonce = nonce24ToWord16 n24, pktData = encrypted }
                                         in return (Just (pkt, pktno+1))
 
+-- | Blocks until a packet is available, or the session is terminated.
+receiveCrypto :: NetCryptoSession -> IO (Maybe CryptoMessage)
+receiveCrypto nc@NCrypto { ncState, ncPacketQueue } = do
+    cd <- atomically $ do
+            st <- readTVar ncState
+            case st of
+                Accepted  -> Just <$> PQ.dequeue ncPacketQueue
+                Confirmed -> Just <$> PQ.dequeue ncPacketQueue
+                _         -> pure Nothing
+    tid <- myThreadId
+    forM_ cd $ runCryptoHook nc
+    return $ bufferData <$> cd
+
 -- | called when we recieve a crypto handshake with valid cookie
 -- TODO set priority on contact addr to 0 if it is older than ForgetPeriod,
 -- then increment it regardless. (Keep addr in MinMaxPSQ in Roster.Contact)
@@ -340,23 +353,24 @@ freshCryptoSession  sessions
                     , ncView = ncView0
                     , ncPacketQueue        = pktq
                     , ncBufferStart        = bufstart
-                    , ncDequeueThread      = Nothing -- error "you want the NetCrypto-Dequeue thread id, but is it started?"
                     , ncPingMachine        = Nothing -- error "you want the NetCrypto-PingMachine, but is it started?"
                     , ncOutgoingQueue      = pktoq
                     }
-    -- launch dequeue thread
-    threadid <- forkIO $ do
-        tid <- myThreadId
-        labelThread tid ("NetCryptoDequeue." ++ show (key2id remotePublicKey))
-        fix $ \loop -> do
-            cd <- atomically $ PQ.dequeue pktq
-            _ <- runCryptoHook (netCryptoSession0 {ncDequeueThread=Just tid}) cd
-            loop
+
+    hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
+
+    -- Dequeue thread:
+    --
+    --   Hopefully, somebody will launch a thread to repeatedly call
+    --   'receiveCrypto' in order to dequeue messages from ncPacketQueue.
+    when (null hooks) $ do
+        hPutStrLn stderr "Warning: Missing new-session handler.  Lost session!"
+
     -- launch ping thread
     fuzz <- randomRIO (0,2000)
     pingMachine <- forkPingMachine ("NetCrypto." ++ show (key2id remotePublicKey)) (15000 + fuzz) 2000
     -- update session with thread ids
-    let netCryptoSession = netCryptoSession0 {ncDequeueThread=Just threadid, ncPingMachine=Just pingMachine}
+    let netCryptoSession = netCryptoSession0 {ncPingMachine=Just pingMachine}
     -- add this session to the lookup maps
     atomically $ do
         modifyTVar allsessions (Map.insert addr netCryptoSession)
@@ -365,7 +379,6 @@ freshCryptoSession  sessions
             Nothing -> modifyTVar allsessionsByKey (Map.insert remotePublicKey [netCryptoSession])
             Just xs -> modifyTVar allsessionsByKey (Map.insert remotePublicKey (netCryptoSession:xs))
     -- run announceNewSessionHooks
-    hooks <- atomically $ readTVar (announceNewSessionHooks sessions)
     flip fix (hooks,netCryptoSession) $ \loop (hooks,session) ->
         case hooks of
             [] -> return ()