1 files changed, 17 insertions, 12 deletions

diff --git a/src/Network/Tox/Onion/Handlers.hs b/src/Network/Tox/Onion/Handlers.hs
index 72398735..08f5cabd 100644
--- a/src/Network/Tox/Onion/Handlers.hs
+++ b/src/Network/Tox/Onion/Handlers.hs
@@ -35,7 +35,7 @@ import Data.Bits
 import Data.Ord
 import Data.Functor.Identity
 
-type Client = QR.Client String PacketKind TransactionId OnionDestination Message
+type Client r = QR.Client String PacketKind TransactionId (OnionDestination r) Message
 type Message = OnionMessage Identity
 
 classify :: Message -> MessageClass String PacketKind TransactionId
@@ -59,7 +59,7 @@ classify msg = go msg
 -- The reason for this 20 second timeout in toxcore is that it gives a reasonable
 -- time (20 to 40 seconds) for a peer to announce himself while taking in count
 -- all the possible delays with some extra seconds.
-announceH :: Routing -> TVar SessionTokens -> TVar AnnouncedKeys -> OnionDestination -> AnnounceRequest -> IO AnnounceResponse
+announceH :: Routing -> TVar SessionTokens -> TVar AnnouncedKeys -> OnionDestination r -> AnnounceRequest -> IO AnnounceResponse
 announceH routing toks keydb (OnionToOwner naddr retpath) req = do
     case () of
         _ | announcePingId req == zeros32
@@ -76,7 +76,7 @@ announceH routing toks keydb (OnionToOwner naddr retpath) req = do
         let storing = (nodeId naddr == announceSeeking req)
         record <- atomically $ do
             when (withTok && storing) $ do
-                let toxpath = OnionToOwner naddr{ nodeId = announceKey req } retpath
+                let toxpath = AnnouncedRoute naddr{ nodeId = announceKey req } retpath
                     -- Note: The following distance calculation assumes that
                     -- our nodeid doesn't change and is the same for both
                     -- routing4 and routing6.
@@ -89,16 +89,16 @@ announceH routing toks keydb (OnionToOwner naddr retpath) req = do
                     then Nonce32 . toPaddedByteString 32 <$> grantToken toks naddr
                     else return $ zeros32
         let k = case record of
-                Nothing                              -> NotStored newtok
-                Just (OnionDestination {}) | storing -> Acknowledged newtok
-                Just (OnionToOwner   ni _)           -> SendBackKey $ id2key (nodeId ni)
+                Nothing                    -> NotStored newtok
+                Just _  | storing          -> Acknowledged newtok
+                Just (AnnouncedRoute ni _) -> SendBackKey $ id2key (nodeId ni)
         let response = AnnounceResponse k ns
         hPutStrLn stderr $ unwords ["Announce:", show req, "-reply->", show response]
         return response
 
 dataToRouteH ::
           TVar AnnouncedKeys
-          -> Transport err OnionDestination (OnionMessage f)
+          -> Transport err (OnionDestination r) (OnionMessage f)
           -> addr
           -> OnionMessage f
           -> IO ()
@@ -111,18 +111,23 @@ dataToRouteH keydb udp _ (OnionToRoute pub assym) = do
             return rpath
     forM_ mb $ \rpath -> do
         -- forward
-        sendMessage udp rpath $ OnionToRouteResponse assym
+        sendMessage udp (toOnionDestination rpath) $ OnionToRouteResponse assym
         hPutStrLn stderr $ "Forwarding data-to-route -->"++show k
 
 type NodeDistance = NodeId
 
+data AnnouncedRoute = AnnouncedRoute NodeInfo (ReturnPath N3)
+
+toOnionDestination :: AnnouncedRoute -> OnionDestination r
+toOnionDestination (AnnouncedRoute ni rpath) = OnionToOwner ni rpath
+
 data AnnouncedKeys = AnnouncedKeys
     { keyByAge :: !(PSQ NodeId (Down POSIXTime)) -- timeout of 300 seconds
-    , keyAssoc :: !(MinMaxPSQ' NodeId NodeDistance (Int,OnionDestination))
+    , keyAssoc :: !(MinMaxPSQ' NodeId NodeDistance (Int,AnnouncedRoute))
     }
 
 
-insertKey :: POSIXTime -> NodeId -> OnionDestination -> NodeDistance -> AnnouncedKeys -> AnnouncedKeys
+insertKey :: POSIXTime -> NodeId -> AnnouncedRoute -> NodeDistance -> AnnouncedKeys -> AnnouncedKeys
 insertKey tm pub toxpath d keydb = AnnouncedKeys
     { keyByAge = PSQ.insert pub (Down tm) (keyByAge keydb)
     , keyAssoc = case MinMaxPSQ.lookup' pub (keyAssoc keydb) of
@@ -134,12 +139,12 @@ areq :: Message -> Either String AnnounceRequest
 areq (OnionAnnounce assym) = Right $ fst $ runIdentity $ assymData assym
 areq _                     = Left "Unexpected non-announce OnionMessage"
 
-handlers :: Transport err OnionDestination Message
+handlers :: Transport err (OnionDestination r) Message
             -> Routing
             -> TVar SessionTokens
             -> TVar AnnouncedKeys
             -> PacketKind
-            -> Maybe (MethodHandler String TransactionId OnionDestination Message)
+            -> Maybe (MethodHandler String TransactionId (OnionDestination r) Message)
 handlers net routing toks keydb AnnounceType
     = Just
     $ MethodHandler areq (\(TransactionId n8 n24) src dst -> OnionAnnounceResponse n8 n24 . Identity)