Option to configure when tunnel is offered

3 files changed, 24 insertions, 1 deletions

diff --git a/main/Main.hs b/main/Main.hs
index 064cf54..403e5e9 100644
--- a/main/Main.hs
+++ b/main/Main.hs
@@ -145,6 +145,11 @@ options =
         (ReqArg (\value -> serviceAttr $ \attrs -> return attrs { discoveryTurnServer = Just (read value) }) "<server>")
         "offer <server> (domain name or IP address) to discovery peers for TURN protocol"
 #endif
+    , Option [] [ "discovery-tunnel" ]
+        (OptArg (\value -> \opts -> do
+            fun <- provideTunnelFun value
+            serviceAttr (\attrs -> return attrs { discoveryProvideTunnel = fun }) opts) "<peer-type>")
+        "offer to provide tunnel for peers of given <peer-type>, possible values: all, none, websocket"
     , Option [] ["dm-bot-echo"]
         (ReqArg (\prefix -> \opts -> return opts { optDmBotEcho = Just (T.pack prefix) }) "<prefix>")
         "automatically reply to direct messages with the same text prefixed with <prefix>"
@@ -173,6 +178,18 @@ 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 (Just name) = do
+        tell [ "Invalid value of --discovery-tunnel: ‘" <> name <> "’\n" ]
+        return $ const False
+
 servicesOptions :: [ OptDescr (Options -> Writer [ String ] Options) ]
 servicesOptions = concatMap helper $ "all" : map soptName availableServices
   where
diff --git a/main/Test.hs b/main/Test.hs
index 62c7229..b07bd87 100644
--- a/main/Test.hs
+++ b/main/Test.hs
@@ -227,6 +227,11 @@ directMessageAttributes out = DirectMessageAttributes
     { dmOwnerMismatch = afterCommit $ outLine out "dm-owner-mismatch"
     }
 
+discoveryAttributes :: DiscoveryAttributes
+discoveryAttributes = (defaultServiceAttributes Proxy)
+    { discoveryProvideTunnel = const True
+    }
+
 dmReceivedWatcher :: Output -> Stored DirectMessage -> IO ()
 dmReceivedWatcher out smsg = do
     let msg = fromStored smsg
@@ -510,7 +515,7 @@ cmdStartServer = do
         "attach" -> return $ someServiceAttr $ pairingAttributes (Proxy @AttachService) out rsPeers "attach"
         "chatroom" -> return $ someService @ChatroomService Proxy
         "contact" -> return $ someServiceAttr $ pairingAttributes (Proxy @ContactService) out rsPeers "contact"
-        "discovery" -> return $ someService @DiscoveryService Proxy
+        "discovery" -> return $ someServiceAttr $ discoveryAttributes
         "dm" -> return $ someServiceAttr $ directMessageAttributes out
         "sync" -> return $ someService @SyncService Proxy
         "test" -> return $ someServiceAttr $ (defaultServiceAttributes Proxy)
diff --git a/main/WebSocket.hs b/main/WebSocket.hs
index fbdd65f..5c49aa1 100644
--- a/main/WebSocket.hs
+++ b/main/WebSocket.hs
@@ -1,4 +1,5 @@
 module WebSocket (
+    WebSocketAddress(..),
     startWebsocketServer,
 ) where