1 files changed, 37 insertions, 2 deletions
diff --git a/src/Erebos/Discovery.hs b/src/Erebos/Discovery.hs
index 0f194a9..3cb55bd 100644
--- a/src/Erebos/Discovery.hs
+++ b/src/Erebos/Discovery.hs
@@ -4,6 +4,8 @@ module Erebos.Discovery (
     DiscoveryService(..),
     DiscoveryAttributes(..),
     DiscoveryConnection(..),
+
+    discoverySearch,
 ) where
 
 import Control.Concurrent
@@ -12,9 +14,13 @@ import Control.Monad.Except
 import Control.Monad.Reader
 
 import Data.IP qualified as IP
+import Data.List
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as M
 import Data.Maybe
+import Data.Proxy
+import Data.Set (Set)
+import Data.Set qualified as S
 import Data.Text (Text)
 import Data.Text qualified as T
 import Data.Word
@@ -165,6 +171,7 @@ data DiscoveryPeerState = DiscoveryPeerState
 
 data DiscoveryGlobalState = DiscoveryGlobalState
     { dgsPeers :: Map RefDigest DiscoveryPeer
+    , dgsSearchingFor :: Set RefDigest
     }
 
 instance Service DiscoveryService where
@@ -181,6 +188,7 @@ instance Service DiscoveryService where
     type ServiceGlobalState DiscoveryService = DiscoveryGlobalState
     emptyServiceGlobalState _ = DiscoveryGlobalState
         { dgsPeers = M.empty
+        , dgsSearchingFor = S.empty
         }
 
     serviceHandler msg = case fromStored msg of
@@ -312,7 +320,7 @@ instance Service DiscoveryService where
         DiscoveryConnectionRequest conn -> do
             self <- svcSelf
             let rconn = emptyConnection (dconnSource conn) (dconnTarget conn)
-            if either refDigest id (dconnTarget conn) `elem` (map (refDigest . storedRef) $ idDataF =<< unfoldOwners self)
+            if either refDigest id (dconnTarget conn) `elem` identityDigests self
               then if
 #ifdef ENABLE_ICE_SUPPORT
                 -- request for us, create ICE sesssion
@@ -347,7 +355,7 @@ instance Service DiscoveryService where
         DiscoveryConnectionResponse conn -> do
             self <- svcSelf
             dpeers <- dgsPeers <$> svcGetGlobal
-            if either refDigest id (dconnSource conn) `elem` (map (refDigest . storedRef) $ idDataF =<< unfoldOwners self)
+            if either refDigest id (dconnSource conn) `elem` identityDigests self
                then do
                     -- response to our request, try to connect to the peer
                     server <- asks svcServer
@@ -389,5 +397,32 @@ instance Service DiscoveryService where
 #endif
             ]
 
+        pid <- asks svcPeerIdentity
+        gs <- svcGetGlobal
+        let searchingFor = foldl' (flip S.delete) (dgsSearchingFor gs) (identityDigests pid)
+        svcModifyGlobal $ \s -> s { dgsSearchingFor = searchingFor }
+
         when (not $ null addrs) $ do
             sendToPeer peer $ DiscoverySelf addrs Nothing
+        forM_ searchingFor $ \dgst -> do
+            sendToPeer peer $ DiscoverySearch (Right dgst)
+
+
+identityDigests :: Foldable f => Identity f -> [ RefDigest ]
+identityDigests pid = map (refDigest . storedRef) $ idDataF =<< unfoldOwners pid
+
+
+discoverySearch :: (MonadIO m, MonadError e m, FromErebosError e) => Server -> RefDigest -> m ()
+discoverySearch server dgst = do
+    peers <- liftIO $ getCurrentPeerList server
+    match <- forM peers $ \peer -> do
+        peerIdentity peer >>= \case
+            PeerIdentityFull pid -> do
+                return $ dgst `elem` identityDigests pid
+            _ -> return False
+    when (not $ or match) $ do
+        modifyServiceGlobalState server (Proxy @DiscoveryService) $ \s -> (, ()) s
+            { dgsSearchingFor = S.insert dgst $ dgsSearchingFor s
+            }
+        forM_ peers $ \peer -> do
+            sendToPeer peer $ DiscoverySearch $ Right dgst