1 files changed, 24 insertions, 14 deletions

diff --git a/src/Network/Lossless.hs b/src/Network/Lossless.hs
index bdbeb3a2..f48dc8fd 100644
--- a/src/Network/Lossless.hs
+++ b/src/Network/Lossless.hs
@@ -27,7 +27,7 @@ data SequenceInfo = SequenceInfo
 
 lossless :: Show addr =>
             (x -> addr -> IO (PacketInboundEvent (x',addr')))
-            -> (SequenceInfo -> x' -> addr' -> IO y)
+            -> (SequenceInfo -> x' -> addr' -> IO (Bool,y))
             -> addr
             -> TransportA String addr x y
             -> IO ( Transport String addr' x'
@@ -39,44 +39,54 @@ lossless isLossless encode saddr udp = do
     oob <- atomically newTChan -- Out-of-band channel, these packets (or
                                -- errors) bypass the packet buffer to be
                                -- received immediately.
-    rloop <- forkIO $ fix $ \loop -> 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
-            forM_ m $ \raw -> do
-            m' <- mapM (uncurry isLossless) raw
+            m' <- mapM (mapM $ uncurry isLossless) m
             case m' of
-                Left e -> do
-                    atomically $ writeTChan oob (Left e)
+                Nothing -> do
+                    atomically $ writeTChan oob Nothing
+                    -- Quit thread here.
+                Just (Left e) -> do
+                    atomically $ writeTChan oob (Just $ Left e)
                     loop
-                Right event -> do
+                Just (Right event) -> do
                     atomically $ do
                         -- x' <- isLossless xaddr x
                         PB.grokInboundPacket pb event
                         case event of
-                            PacketReceivedLossy {} -> writeTChan oob (Right $ peReceivedPayload event)
+                            PacketReceivedLossy {} -> writeTChan oob (Just $ Right $ peReceivedPayload event)
                             _                      -> do
                                 report <- pbReport "enqueued" pb
-                                writeTChan oob (Left report)
+                                writeTChan oob (Just $ Left report)
                     loop
     let tr = Transport
             { awaitMessage = \kont -> do
                 join $ atomically $ orElse
                     (do x <- readTChan oob
-                        return $ kont $! Just x)
+                        return $ kont $! x)
                     (do x <- PB.awaitReadyPacket pb
                         report <- pbReport "dequeued" pb
                         return $ do
-                            dput XNetCrypto report
+                            atomically $ writeTChan oob (Just $ Left report)
                             kont $! Just (Right x))
             , sendMessage = \a' x' -> do
                 seqno <- atomically $ do
                     seqno <- PB.nextToSendSequenceNumber pb
                     ack   <- PB.expectingSequenceNumber pb
                     return $ SequenceInfo seqno ack
-                x <- encode seqno x' a'
-                (isfull,nn) <- atomically $ PB.grokOutboundPacket pb (PacketSent (sequenceNumber seqno) x)
+                (islossy,x) <- encode seqno x' a'
+                (isfull,nn) <-
+                    if islossy
+                        then do
+                            dput XNetCrypto $ shows saddr $ " <-- Lossy packet " ++ show seqno
+                            return (False,(0,0)) -- avoid updating seqno on lossy packets.
+                        else do
+                            dput XNetCrypto $ shows saddr $ " <-- Lossless packet " ++ show seqno
+                            atomically $ PB.grokOutboundPacket pb (PacketSent (sequenceNumber seqno) x)
                 when isfull $ do
                     dput XNetCrypto $ shows saddr $ " <-- Outbound queue is full! Retrying... " ++ show (nn,seqno)
                     atomically $ do
@@ -84,7 +94,7 @@ lossless isLossless encode saddr udp = do
                         when isfull retry
                 sendMessage udp saddr x
             , closeTransport = do
-                killThread rloop
+                atomically $ writeTChan oob Nothing -- quit rloop thread
                 closeTransport udp
             }
         resend ns = do