STM-based awaitMessage.

1 files changed, 8 insertions, 7 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