Clean up forked transports.

1 files changed, 6 insertions, 6 deletions

diff --git a/dht/src/Network/QueryResponse.hs b/dht/src/Network/QueryResponse.hs
index 4f956936..c7ab59d8 100644
--- a/dht/src/Network/QueryResponse.hs
+++ b/dht/src/Network/QueryResponse.hs
@@ -124,15 +124,15 @@ partitionTransportM parse encodex tr = do
                     awaitMessage tr $ \m -> case m of
                         Just (Right msg) -> parse msg >>=
                                                 either (kont . Just . Right)
-                                                       (\y -> putMVar mvar y >> again)
+                                                       (\y -> putMVar mvar (Just y) >> again)
                         Just (Left e)    -> kont $ Just (Left e)
-                        Nothing          -> kont Nothing
+                        Nothing          -> putMVar mvar Nothing >> kont Nothing
                  , sendMessage = \addr' msg' -> do
                     msg_addr <- encodex (msg',addr')
                     mapM_ (uncurry . flip $ sendMessage tr) msg_addr
                  }
         ytr = Transport
-                { awaitMessage = \kont -> takeMVar mvar >>= kont . Just . Right
+                { awaitMessage = \kont -> takeMVar mvar >>= kont . fmap Right
                 , sendMessage = sendMessage tr
                 , closeTransport = return ()
                 }
@@ -150,9 +150,9 @@ partitionAndForkTransport forkedSend parse encodex tr = do
                     awaitMessage tr $ \m -> case m of
                         Just (Right msg) -> parse msg >>=
                                                 either (kont . Just . Right)
-                                                       (\y -> putMVar mvar y >> again)
+                                                       (\y -> putMVar mvar (Just y) >> again)
                         Just (Left e)    -> kont $ Just (Left e)
-                        Nothing          -> kont Nothing
+                        Nothing          -> putMVar mvar Nothing >> kont Nothing
                  , sendMessage = \addr' msg' -> do
                     msg_addr <- encodex (msg',addr')
                     case msg_addr of
@@ -161,7 +161,7 @@ partitionAndForkTransport forkedSend parse encodex tr = do
                         Nothing               -> return ()
                  }
         ytr = Transport
-                { awaitMessage   = \kont -> takeMVar mvar >>= kont . Just . Right
+                { awaitMessage   = \kont -> takeMVar mvar >>= kont . fmap Right
                 , sendMessage    = sendMessage tr
                 , closeTransport = return ()
                 }