Peer address input in ServiceHandler

4 files changed, 11 insertions, 6 deletions

diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index 3c15f94..ff07964 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -66,7 +66,7 @@ data DiscoveryAttributes = DiscoveryAttributes
     , discoveryStunServer :: Maybe Text
     , discoveryTurnPort :: Maybe Word16
     , discoveryTurnServer :: Maybe Text
-    , discoveryProvideTunnel :: Peer -> Bool
+    , discoveryProvideTunnel :: Peer -> PeerAddress -> Bool
     }
 
 defaultDiscoveryAttributes :: DiscoveryAttributes
@@ -75,7 +75,7 @@ defaultDiscoveryAttributes = DiscoveryAttributes
     , discoveryStunServer = Nothing
     , discoveryTurnPort = Nothing
     , discoveryTurnServer = Nothing
-    , discoveryProvideTunnel = const False
+    , discoveryProvideTunnel = \_ _ -> False
     }
 
 data DiscoveryConnection = DiscoveryConnection
@@ -276,7 +276,7 @@ instance Service DiscoveryService where
                 (discoveryTurnPort attrs)
 
         DiscoveryAcknowledged _ stunServer stunPort turnServer turnPort -> do
-            paddr <- asks (peerAddress . svcPeer) >>= return . \case
+            paddr <- asks svcPeerAddress >>= return . \case
                 (DatagramAddress saddr) -> T.pack . show . fst <$> inetFromSockAddr saddr
                 _ -> Nothing
 
@@ -293,10 +293,11 @@ instance Service DiscoveryService where
         DiscoverySearch edgst -> do
             dpeer <- M.lookup (either refDigest id edgst) . dgsPeers <$> svcGetGlobal
             peer <- asks svcPeer
+            paddr <- asks svcPeerAddress
             attrs <- asks svcAttributes
             let offerTunnel
-                    | discoveryProvideTunnel attrs peer = (++ [ DiscoveryTunnel ])
-                    | otherwise                         = id
+                    | discoveryProvideTunnel attrs peer paddr = (++ [ DiscoveryTunnel ])
+                    | otherwise                               = id
             replyPacket $ DiscoveryResult edgst $ maybe [] (offerTunnel . dpAddress) dpeer
 
         DiscoveryResult _ [] -> do
@@ -409,6 +410,7 @@ instance Service DiscoveryService where
               else do
                 -- request to some of our peers, relay
                 peer <- asks svcPeer
+                paddr <- asks svcPeerAddress
                 mbdp <- M.lookup (either refDigest id $ dconnTarget conn) . dgsPeers <$> svcGetGlobal
                 streams <- receivedStreams
                 case mbdp of
@@ -416,7 +418,7 @@ instance Service DiscoveryService where
                         Just dp
                             | Just dpeer <- dpPeer dp -> if
                                 | dconnTunnel conn -> if
-                                    | not (discoveryProvideTunnel attrs peer) -> do
+                                    | not (discoveryProvideTunnel attrs peer paddr) -> do
                                         replyPacket $ DiscoveryConnectionResponse rconn
                                     | fromSource : _ <- streams -> do
                                         void $ liftIO $ forkIO $ runPeerService @DiscoveryService dpeer $ do
diff --git a/src/Erebos/Network.hs b/src/Erebos/Network.hs
index 3aa1a3f..76ecd82 100644
--- a/src/Erebos/Network.hs
+++ b/src/Erebos/Network.hs
@@ -1001,6 +1001,7 @@ runPeerServiceOn mbservice newStreams peer handler = liftIO $ do
                             let inp = ServiceInput
                                     { svcAttributes = attr
                                     , svcPeer = peer
+                                    , svcPeerAddress = peerAddress peer
                                     , svcPeerIdentity = peerId
                                     , svcServer = server
                                     , svcPrintOp = atomically . logd
diff --git a/src/Erebos/Network.hs-boot b/src/Erebos/Network.hs-boot
index af77581..17a5275 100644
--- a/src/Erebos/Network.hs-boot
+++ b/src/Erebos/Network.hs-boot
@@ -4,5 +4,6 @@ import Erebos.Object.Internal
 
 data Server
 data Peer
+data PeerAddress
 
 peerStorage :: Peer -> Storage
diff --git a/src/Erebos/Service.hs b/src/Erebos/Service.hs
index 4499ef9..afcf512 100644
--- a/src/Erebos/Service.hs
+++ b/src/Erebos/Service.hs
@@ -115,6 +115,7 @@ mkServiceID = maybe (error "Invalid service ID") ServiceID . U.fromString
 data ServiceInput s = ServiceInput
     { svcAttributes :: ServiceAttributes s
     , svcPeer :: Peer
+    , svcPeerAddress :: PeerAddress
     , svcPeerIdentity :: UnifiedIdentity
     , svcServer :: Server
     , svcPrintOp :: String -> IO ()