Discovery on localhost

Changelog: Discovery is now working also on localhost without need for external interface.

3 files changed, 149 insertions, 19 deletions

diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index 9f25b37..da55dc3 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -60,7 +60,8 @@ data DiscoveryService
     | DiscoveryConnectionResponse DiscoveryConnection
 
 data DiscoveryAddress
-    = DiscoveryIP InetAddress PortNumber
+    = DiscoveryLocalhost PortNumber
+    | DiscoveryIP InetAddress PortNumber
     | DiscoveryICE
     | DiscoveryTunnel
     | DiscoveryOther Text
@@ -185,6 +186,7 @@ instance Storable DiscoveryService where
 
 instance StorableText DiscoveryAddress where
     toText = \case
+        DiscoveryLocalhost port -> T.unwords [ "localhost", T.pack $ show port ]
         DiscoveryIP addr port -> T.unwords [ T.pack $ show addr, T.pack $ show port ]
         DiscoveryICE -> "ICE"
         DiscoveryTunnel -> "tunnel"
@@ -192,6 +194,11 @@ instance StorableText DiscoveryAddress where
 
     fromText str = return $ if
         | [ addrStr, portStr ] <- T.words str
+        , addrStr == "localhost"
+        , Just port <- readMaybe $ T.unpack portStr
+        -> DiscoveryLocalhost port
+
+        | [ addrStr, portStr ] <- T.words str
         , Just addr <- readMaybe $ T.unpack addrStr
         , Just port <- readMaybe $ T.unpack portStr
         -> DiscoveryIP addr port
@@ -219,6 +226,7 @@ data DiscoveryPeer = DiscoveryPeer
     { dpPriority :: Int
     , dpPeer :: Maybe Peer
     , dpAddress :: [ DiscoveryAddress ]
+    , dpLocalhost :: Maybe PortNumber
     , dpIceSession :: Maybe IceSession
     }
 
@@ -227,6 +235,7 @@ emptyPeer = DiscoveryPeer
     { dpPriority = 0
     , dpPeer = Nothing
     , dpAddress = []
+    , dpLocalhost = Nothing
     , dpIceSession = Nothing
     }
 
@@ -314,24 +323,28 @@ instance Service DiscoveryService where
             peer <- asks svcPeer
             paddrs <- getPeerAddresses peer
 
+            let matchedAddrs = flip filter addrs $ \case
+                    DiscoveryICE -> True
+                    DiscoveryIP ipaddr port ->
+                        DatagramAddress (inetToSockAddr ( ipaddr, port )) `elem` paddrs
+                    _ -> False
+                localhostPort = getLocalhostPortFromPeerAddresses paddrs
+
             debugLog $ unwords
                 [ "new peer"
                 , show [ refDigest $ storedRef $ idData pid, refDigest $ storedRef $ idExtData pid ]
                 , show $ map (refDigest . storedRef) $ idDataF $ finalOwner pid
                 , show paddrs
+                , show (map toText matchedAddrs)
+                , show localhostPort
                 ]
 
-            let matchedAddrs = flip filter addrs $ \case
-                    DiscoveryICE -> True
-                    DiscoveryIP ipaddr port ->
-                        DatagramAddress (inetToSockAddr ( ipaddr, port )) `elem` paddrs
-                    _ -> False
-
             forM_ (idDataF =<< unfoldOwners pid) $ \sdata -> do
                 let dp = DiscoveryPeer
                         { dpPriority = fromMaybe 0 priority
                         , dpPeer = Just peer
                         , dpAddress = matchedAddrs
+                        , dpLocalhost = localhostPort
                         , dpIceSession = Nothing
                         }
                     rv = ResultValue
@@ -360,10 +373,19 @@ instance Service DiscoveryService where
                             st <- getStorage
                             forM_ dgsts $ \dgst -> do
                                 when (dgst `M.member` peerSearchingFor) $ do
-                                    offerTunnel <- offerTunnelBetween attrs peer sp >>= return . \case
-                                        True  -> (++ [ DiscoveryTunnel ])
-                                        False -> id
-                                    let discoveryAddrs = offerTunnel matchedAddrs
+                                    discoveryAddrs <- concat <$> sequence
+                                        [ case localhostPort of
+                                              Just port -> getPeerAddresses sp >>= \case
+                                                  spaddrs
+                                                      | isJust $ getLocalhostPortFromPeerAddresses spaddrs
+                                                      -> return [ DiscoveryLocalhost port ]
+                                                  _ -> return []
+                                              _ -> return []
+                                        , return matchedAddrs
+                                        , offerTunnelBetween attrs peer sp >>= return . \case
+                                            True  -> [ DiscoveryTunnel ]
+                                            False -> []
+                                        ]
                                     let ( results, via )
                                             | dgst == (refDigest $ storedRef $ idData pid)
                                             = ( discoveryAddrs, [] )
@@ -403,6 +425,7 @@ instance Service DiscoveryService where
             let dgst = either refDigest id edgst
             peer <- asks svcPeer
             pid <- asks svcPeerIdentity
+            paddrs <- getPeerAddresses peer
             (M.lookup dgst . dgsPeers <$> svcGetGlobal) >>= \case
                 Just rv
                   | direct <- (\p -> Just p <* guard (p /= peer)) =<< dpPeer =<< rvDirect rv
@@ -412,14 +435,19 @@ instance Service DiscoveryService where
                   , isJust direct || not (null rvia)
                   -> do
                     attrs <- asks svcAttributes
-                    offerTunnel <- case direct of
-                        Just dpeer -> offerTunnelBetween attrs peer dpeer >>= return . \case
-                            True  -> (++ [ DiscoveryTunnel ])
-                            False -> id
-                        Nothing -> return id
-                    let results
-                            | isJust direct = offerTunnel $ maybe [] dpAddress $ rvDirect rv
-                            | otherwise = []
+                    results <- if
+                        | isJust direct -> concat <$> sequence
+                            [ case ( dpLocalhost =<< rvDirect rv, getLocalhostPortFromPeerAddresses paddrs ) of
+                                ( Just port, Just _ ) -> return [ DiscoveryLocalhost port ]
+                                _ -> return []
+                            , return $ maybe [] dpAddress $ rvDirect rv
+                            , case direct of
+                                Just dpeer -> offerTunnelBetween attrs peer dpeer >>= return . \case
+                                    True  -> [ DiscoveryTunnel ]
+                                    False -> []
+                                Nothing -> return []
+                            ]
+                        | otherwise -> return []
                     via <- liftIO $ fmap catMaybes $ mapM (viaFromPeer attrs peer) rvia
                     replyPacket $ DiscoveryResult edgst results via
                     debugLog $ "search by " <> show (refDigest $ storedRef $ idData pid) <>
@@ -464,6 +492,14 @@ instance Service DiscoveryService where
             let runAsService = runPeerService @DiscoveryService @IO discoveryPeer
 
             let tryAddresses = \case
+                    DiscoveryLocalhost port : _ -> do
+                        void $ liftIO $ forkIO $ do
+                            let saddr = makeLocalhostAddress port
+                            peer <- serverPeer server saddr
+                            runAsService $ do
+                                let upd rv = rv { rvDirect = Just $ (fromMaybe emptyPeer $ rvDirect rv) { dpPeer = Just peer } }
+                                svcModifyGlobal $ \s -> s { dgsPeers = M.alter (Just . upd . fromMaybe mempty) dgst $ dgsPeers s }
+
                     DiscoveryIP ipaddr port : _ -> do
                         void $ liftIO $ forkIO $ do
                             let saddr = inetToSockAddr ( ipaddr, port )
@@ -883,3 +919,6 @@ discoverySetupTunnelResponse target = do
             (emptyConnection (Right self) (Right target))
             { dconnTunnel = True
             }
+
+getLocalhostPortFromPeerAddresses :: [ PeerAddress ] -> Maybe PortNumber
+getLocalhostPortFromPeerAddresses = msum . map (\case DatagramAddress a -> getLocalhostPort a; _ -> Nothing)
diff --git a/src/Erebos/Network/Address.hs b/src/Erebos/Network/Address.hs
index 63f6af1..8dedecd 100644
--- a/src/Erebos/Network/Address.hs
+++ b/src/Erebos/Network/Address.hs
@@ -2,6 +2,8 @@ module Erebos.Network.Address (
     InetAddress(..),
     inetFromSockAddr,
     inetToSockAddr,
+    makeLocalhostAddress,
+    getLocalhostPort,
 
     SockAddr, PortNumber,
 ) where
@@ -63,3 +65,11 @@ inetFromSockAddr saddr = first InetAddress <$> IP.fromSockAddr saddr
 
 inetToSockAddr :: ( InetAddress, PortNumber ) -> SockAddr
 inetToSockAddr = IP.toSockAddr . first fromInetAddress
+
+
+makeLocalhostAddress :: PortNumber -> SockAddr
+makeLocalhostAddress port = SockAddrInet6 port 0 ( 0, 0, 0, 1 ) 0
+
+getLocalhostPort :: SockAddr -> Maybe PortNumber
+getLocalhostPort (SockAddrInet6 port _ ( 0, 0, 0, 1 ) _) = Just port
+getLocalhostPort _                                       = Nothing
diff --git a/test/discovery.et b/test/discovery.et
index b8d5c81..46a5f10 100644
--- a/test/discovery.et
+++ b/test/discovery.et
@@ -487,3 +487,84 @@ test CommonOwnerDiscovery:
         expect from p1:
             /peer [0-9]+ addr ${p2.node.ip} 29665/
             /peer [0-9]+ addr ${p3.node.ip} 29665/
+
+
+test DiscoveryOnLocalhost:
+    let services = "discovery"
+    node n
+
+    # TODO: some way to have node without any connection
+    shell on n:
+        ip link set dev ${n.ifname} down
+
+    spawn on n:
+        as pd
+        as ps # connection order to pd:
+        as p1 # search - p1 - ps
+        as p2 # search - ps - p2
+        as p3 # p3 - search - ps
+        as p4 # p4 - ps - search
+        as p5 # ps - p5 - search
+        as p6 # ps - search - p6
+
+    send "create-identity Discovery" to pd
+    send "create-identity Searching" to ps
+    send "create-identity Device1" to p1
+    send "create-identity Device2" to p2
+    send "create-identity Device3" to p3
+    send "create-identity Device4" to p4
+    send "create-identity Device5" to p5
+    send "create-identity Device6" to p6
+
+    expect /create-identity-done ref $refpat.*/ from pd
+    expect /create-identity-done ref ($refpat).*/ from ps capture ps_id
+    expect /create-identity-done ref ($refpat).*/ from p1 capture p1_id
+    expect /create-identity-done ref ($refpat).*/ from p2 capture p2_id
+    expect /create-identity-done ref ($refpat).*/ from p3 capture p3_id
+    expect /create-identity-done ref ($refpat).*/ from p4 capture p4_id
+    expect /create-identity-done ref ($refpat).*/ from p5 capture p5_id
+    expect /create-identity-done ref ($refpat).*/ from p6 capture p6_id
+
+    send "start-server services $services" to pd
+    # wait for the first server to start and claim the discovery port
+    expect /start-server-done/ from pd
+
+    for p in [ ps, p1, p2, p3, p4, p5, p6 ]:
+        send "start-server services $services" to p
+    for p in [ ps, p1, p2, p3, p4, p5, p6 ]:
+        expect /start-server-done/ from p
+
+    send "discovery-connect $p1_id" to ps
+    send "discovery-connect $p2_id" to ps
+
+    send "peer-add localhost" to p1
+    send "peer-add localhost" to p3
+    send "peer-add localhost" to p4
+    expect /peer [0-9]+ id Discovery/ from p1
+    expect /peer [0-9]+ id Discovery/ from p3
+    expect /peer [0-9]+ id Discovery/ from p4
+
+    send "discovery-connect $p3_id" to ps
+
+    send "peer-add localhost" to ps
+    expect /peer [0-9]+ id Discovery/ from ps
+
+    expect /peer [0-9]+ id Device1/ from ps
+    expect /peer [0-9]+ id Device3/ from ps
+
+    send "peer-add localhost" to p2
+    expect /peer [0-9]+ id Device2/ from ps
+
+    send "discovery-connect $p4_id" to ps
+    expect /peer [0-9]+ id Device4/ from ps
+
+    send "peer-add localhost" to p5
+    expect /peer [0-9]+ id Discovery/ from p5
+
+    send "discovery-connect $p5_id" to ps
+    send "discovery-connect $p6_id" to ps
+    expect /peer [0-9]+ id Device5/ from ps
+
+    send "peer-add localhost" to p6
+    expect /peer [0-9]+ id Discovery/ from p6
+    expect /peer [0-9]+ id Device6/ from ps