Peer address input in ServiceHandler

6 files changed, 21 insertions, 17 deletions

diff --git a/main/Main.hs b/main/Main.hs
index e2e585a..6a6c565 100644
--- a/main/Main.hs
+++ b/main/Main.hs
@@ -172,17 +172,16 @@ options =
             return sopt { soptService = service }
         return opts { optServices = services' }
 
-    provideTunnelFun :: Maybe String -> Writer [ String ] (Peer -> Bool)
-    provideTunnelFun Nothing = return $ const True
-    provideTunnelFun (Just "all") = return $ const True
-    provideTunnelFun (Just "none") = return $ const False
-    provideTunnelFun (Just "websocket") = return $ \peer ->
-        case peerAddress peer of
-            CustomPeerAddress addr | Just WebSocketAddress {} <- cast addr -> True
-            _ -> False
+    provideTunnelFun :: Maybe String -> Writer [ String ] (Peer -> PeerAddress -> Bool)
+    provideTunnelFun Nothing = return $ \_ _ -> True
+    provideTunnelFun (Just "all") = return $ \_ _ -> True
+    provideTunnelFun (Just "none") = return $ \_ _ -> False
+    provideTunnelFun (Just "websocket") = return $ \_ -> \case
+        CustomPeerAddress addr | Just WebSocketAddress {} <- cast addr -> True
+        _ -> False
     provideTunnelFun (Just name) = do
         tell [ "Invalid value of --discovery-tunnel: ‘" <> name <> "’\n" ]
-        return $ const False
+        return $ \_ _ -> False
 
 servicesOptions :: [ OptDescr (Options -> Writer [ String ] Options) ]
 servicesOptions = concatMap helper $ "all" : map soptName availableServices
diff --git a/main/Test.hs b/main/Test.hs
index b1c8f01..093d3ac 100644
--- a/main/Test.hs
+++ b/main/Test.hs
@@ -229,7 +229,7 @@ directMessageAttributes out = DirectMessageAttributes
 
 discoveryAttributes :: DiscoveryAttributes
 discoveryAttributes = (defaultServiceAttributes Proxy)
-    { discoveryProvideTunnel = const False
+    { discoveryProvideTunnel = \_ _ -> False
     }
 
 dmReceivedWatcher :: Output -> Stored DirectMessage -> IO ()
@@ -524,7 +524,7 @@ cmdStartServer = do
         ( "chatroom", _ ) -> return $ someService @ChatroomService Proxy
         ( "contact", _ ) -> return $ someServiceAttr $ pairingAttributes (Proxy @ContactService) out rsPeers "contact"
         ( "discovery", params ) -> return $ someServiceAttr $ discoveryAttributes
-            { discoveryProvideTunnel = const $ "tunnel" `elem` params
+            { discoveryProvideTunnel = \_ _ -> "tunnel" `elem` params
             }
         ( "dm", _ ) -> return $ someServiceAttr $ directMessageAttributes out
         ( "sync", _ ) -> return $ someService @SyncService Proxy
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 ()