hashtree_cli/nostrdb_integration/

crawler.rs

//! Social graph crawler - BFS crawl of follow lists via Nostr relays.

use nostrdb_social::Ndb;
use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::watch;

/// Crawls the social graph by fetching kind 3 (contact list) events from relays
/// and ingesting them into nostrdb.
pub struct SocialGraphCrawler {
    ndb: Arc<Ndb>,
    spambox: Option<Arc<Ndb>>,
    keys: nostr::Keys,
    relays: Vec<String>,
    max_depth: u32,
}

impl SocialGraphCrawler {
    pub fn new(ndb: Arc<Ndb>, keys: nostr::Keys, relays: Vec<String>, max_depth: u32) -> Self {
        Self {
            ndb,
            spambox: None,
            keys,
            relays,
            max_depth,
        }
    }

    pub fn with_spambox(mut self, spambox: Arc<Ndb>) -> Self {
        self.spambox = Some(spambox);
        self
    }

    fn is_within_social_graph(&self, pk_bytes: &[u8; 32]) -> bool {
        if pk_bytes == &self.keys.public_key().to_bytes() {
            return true;
        }
        super::get_follow_distance(&self.ndb, pk_bytes)
            .map(|distance| distance <= self.max_depth)
            .unwrap_or(false)
    }

    fn ingest_event_into(&self, ndb: &Ndb, sub_id: &str, event: &nostr::Event) {
        if let Ok(json) = serde_json::to_string(event) {
            super::ingest_event(ndb, sub_id, &json);
        }
    }

    #[allow(deprecated)] // nostr 0.35 deprecates tags() but we use this version
    fn collect_missing_root_follows(
        &self,
        event: &nostr::Event,
        fetched_contact_lists: &mut HashSet<[u8; 32]>,
    ) -> Vec<[u8; 32]> {
        if self.max_depth < 2 {
            return Vec::new();
        }
        if event.kind != nostr::Kind::ContactList {
            return Vec::new();
        }

        let root_pk = self.keys.public_key().to_bytes();
        if event.pubkey.to_bytes() != root_pk {
            return Vec::new();
        }

        let mut missing = Vec::new();
        for tag in event.tags().iter() {
            if let Some(nostr::TagStandard::PublicKey { public_key, .. }) = tag.as_standardized() {
                let pk_bytes = public_key.to_bytes();
                if fetched_contact_lists.contains(&pk_bytes) {
                    continue;
                }

                let existing_follows = super::get_follows(&self.ndb, &pk_bytes);
                if !existing_follows.is_empty() {
                    fetched_contact_lists.insert(pk_bytes);
                    continue;
                }

                fetched_contact_lists.insert(pk_bytes);
                missing.push(pk_bytes);
            }
        }

        missing
    }

    async fn fetch_contact_lists_for_pubkeys(
        &self,
        client: &nostr_sdk::Client,
        pubkeys: &[[u8; 32]],
        shutdown_rx: &watch::Receiver<bool>,
        sub_id: &str,
    ) {
        for pk_bytes in pubkeys {
            if *shutdown_rx.borrow() {
                break;
            }

            let pk = match nostr::PublicKey::from_slice(pk_bytes) {
                Ok(pk) => pk,
                Err(_) => continue,
            };

            let filter = nostr::Filter::new()
                .author(pk)
                .kinds(vec![nostr::Kind::ContactList, nostr::Kind::Custom(10000)]);

            let source = nostr_sdk::EventSource::relays(Some(Duration::from_secs(5)));
            match tokio::time::timeout(
                Duration::from_secs(10),
                client.get_events_of(vec![filter], source),
            )
            .await
            {
                Ok(Ok(events)) => {
                    for event in &events {
                        self.ingest_event_into(&self.ndb, sub_id, event);
                    }
                }
                Ok(Err(e)) => {
                    tracing::debug!("Failed to fetch events for {}...: {}", &pk.to_hex()[..8], e);
                }
                Err(_) => {
                    tracing::debug!("Timeout fetching events for {}...", &pk.to_hex()[..8]);
                }
            }
        }
    }

    pub(crate) fn handle_incoming_event(&self, event: &nostr::Event) {
        let is_contact_list = event.kind == nostr::Kind::ContactList;
        let is_mute_list = event.kind == nostr::Kind::Custom(10000);
        if !is_contact_list && !is_mute_list {
            return;
        }

        let pk_bytes = event.pubkey.to_bytes();
        if self.is_within_social_graph(&pk_bytes) {
            self.ingest_event_into(&self.ndb, "live", event);
            return;
        }

        if let Some(spambox) = &self.spambox {
            self.ingest_event_into(spambox, "spambox", event);
        } else {
            tracing::debug!(
                "Social graph crawler: dropping untrusted {} from {}...",
                if is_contact_list {
                    "contact list"
                } else {
                    "mute list"
                },
                &event.pubkey.to_hex()[..8.min(event.pubkey.to_hex().len())]
            );
        }
    }

    /// Run the BFS crawl until shutdown is signaled.
    /// Fetches contact lists from relays and feeds them into nostrdb.
    #[allow(deprecated)] // nostr 0.35 deprecates kind()/tags() but we use this version
    pub async fn crawl(&self, shutdown_rx: watch::Receiver<bool>) {
        use nostr::nips::nip19::ToBech32;
        use nostr_sdk::prelude::RelayPoolNotification;

        if self.relays.is_empty() {
            tracing::warn!("Social graph crawler: no relays configured, skipping");
            return;
        }

        let mut shutdown_rx = shutdown_rx;
        if *shutdown_rx.borrow() {
            tracing::info!("Social graph crawler: shutdown requested before start");
            return;
        }

        tracing::info!(
            "Starting social graph crawl (max_depth={}, relays={})",
            self.max_depth,
            self.relays.len()
        );

        let sdk_keys =
            match nostr_sdk::Keys::parse(&self.keys.secret_key().to_bech32().unwrap_or_default()) {
                Ok(k) => k,
                Err(e) => {
                    tracing::error!("Failed to parse keys for crawler: {}", e);
                    return;
                }
            };

        let client = nostr_sdk::Client::new(&sdk_keys);

        for relay in &self.relays {
            if let Err(e) = client.add_relay(relay).await {
                tracing::warn!("Failed to add relay {}: {}", relay, e);
            }
        }
        client.connect().await;

        // BFS: start from root, fetch contact lists at each depth
        let root_pk = self.keys.public_key().to_bytes();
        let mut visited: HashSet<[u8; 32]> = HashSet::new();
        let mut fetched_contact_lists: HashSet<[u8; 32]> = HashSet::new();
        let mut current_level = vec![root_pk];
        visited.insert(root_pk);

        for depth in 0..self.max_depth {
            if current_level.is_empty() || *shutdown_rx.borrow() {
                break;
            }

            tracing::info!(
                "Crawling depth {} with {} pubkeys",
                depth,
                current_level.len()
            );

            let mut next_level = Vec::new();

            for pk_bytes in &current_level {
                if *shutdown_rx.borrow() {
                    break;
                }

                fetched_contact_lists.insert(*pk_bytes);

                let pk_hex = hex::encode(pk_bytes);

                let pk = match nostr::PublicKey::from_slice(pk_bytes) {
                    Ok(pk) => pk,
                    Err(_) => continue,
                };

                let filter = nostr::Filter::new()
                    .author(pk)
                    .kinds(vec![nostr::Kind::ContactList, nostr::Kind::Custom(10000)]);

                let source = nostr_sdk::EventSource::relays(Some(Duration::from_secs(5)));

                match tokio::time::timeout(
                    Duration::from_secs(10),
                    client.get_events_of(vec![filter], source),
                )
                .await
                {
                    Ok(Ok(events)) => {
                        for event in &events {
                            // Ingest into nostrdb
                            self.ingest_event_into(&self.ndb, "crawl", event);

                            // Extract follows for next level
                            if event.kind() == nostr::Kind::ContactList {
                                for tag in event.tags().iter() {
                                    if let Some(nostr::TagStandard::PublicKey {
                                        public_key, ..
                                    }) = tag.as_standardized()
                                    {
                                        let follow_bytes = public_key.to_bytes();
                                        if !visited.contains(&follow_bytes) {
                                            visited.insert(follow_bytes);
                                            next_level.push(follow_bytes);
                                        }
                                    }
                                }
                            }
                        }
                        tracing::debug!(
                            "Depth {}: fetched {} events for {}...",
                            depth,
                            events.len(),
                            &pk_hex[..8.min(pk_hex.len())]
                        );
                    }
                    Ok(Err(e)) => {
                        tracing::debug!("Failed to fetch events for {}...: {}", &pk_hex[..8], e);
                    }
                    Err(_) => {
                        tracing::debug!("Timeout fetching events for {}...", &pk_hex[..8]);
                    }
                }
            }

            current_level = next_level;
        }

        let filter = nostr::Filter::new()
            .kinds(vec![nostr::Kind::ContactList, nostr::Kind::Custom(10000)])
            .since(nostr::Timestamp::now());

        match client.subscribe(vec![filter], None).await {
            Ok(_) => tracing::info!("Social graph crawler: subscribed to contact and mute lists"),
            Err(e) => tracing::warn!("Social graph crawler: failed to subscribe: {}", e),
        }

        let mut notifications = client.notifications();
        loop {
            tokio::select! {
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        break;
                    }
                }
                notification = notifications.recv() => {
                    match notification {
                        Ok(RelayPoolNotification::Event { event, .. }) => {
                            self.handle_incoming_event(&event);

                            let missing = self.collect_missing_root_follows(&event, &mut fetched_contact_lists);
                            if !missing.is_empty() {
                                tracing::info!(
                                    "Root follow list updated: fetching {} missing contact lists",
                                    missing.len()
                                );
                                self.fetch_contact_lists_for_pubkeys(&client, &missing, &shutdown_rx, "recrawl").await;
                            }
                        }
                        Ok(_) => {}
                        Err(e) => {
                            tracing::warn!("Social graph crawler notification error: {}", e);
                            break;
                        }
                    }
                }
            }
        }

        if let Err(e) = client.disconnect().await {
            tracing::debug!("Error disconnecting crawler client: {}", e);
        }

        tracing::info!(
            "Social graph crawl complete: visited {} pubkeys",
            visited.len()
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nostr::{EventBuilder, Kind, PublicKey, Tag};
    use std::collections::HashSet;
    use std::sync::Arc;
    use tempfile::TempDir;

    async fn wait_for_follow(ndb: &Ndb, owner: &[u8; 32], target: &[u8; 32]) -> bool {
        let deadline = tokio::time::Instant::now() + Duration::from_millis(500);
        loop {
            let follows = super::super::get_follows(ndb, owner);
            if follows.iter().any(|pk| pk == target) {
                return true;
            }
            if tokio::time::Instant::now() >= deadline {
                return false;
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    }

    #[tokio::test]
    async fn test_crawler_routes_untrusted_to_spambox() {
        let _guard = super::super::test_lock();
        let tmp = TempDir::new().unwrap();
        let ndb = super::super::init_ndb(tmp.path()).unwrap();
        let spambox =
            super::super::init_ndb_at_path(&tmp.path().join("nostrdb_spambox"), None).unwrap();

        let root_keys = nostr::Keys::generate();
        let root_pk = root_keys.public_key().to_bytes();
        super::super::set_social_graph_root(&ndb, &root_pk);
        tokio::time::sleep(Duration::from_millis(100)).await;

        let crawler = SocialGraphCrawler::new(Arc::clone(&ndb), root_keys.clone(), vec![], 2)
            .with_spambox(Arc::clone(&spambox));

        let unknown_keys = nostr::Keys::generate();
        let follow_tag = Tag::public_key(PublicKey::from_slice(&root_pk).unwrap());
        let event = EventBuilder::new(Kind::ContactList, "", vec![follow_tag])
            .to_event(&unknown_keys)
            .unwrap();

        crawler.handle_incoming_event(&event);

        let unknown_pk = unknown_keys.public_key().to_bytes();
        assert!(!wait_for_follow(&ndb, &unknown_pk, &root_pk).await);
        assert!(wait_for_follow(&spambox, &unknown_pk, &root_pk).await);
    }

    #[tokio::test]
    async fn test_crawler_routes_trusted_to_main_db() {
        let _guard = super::super::test_lock();
        let tmp = TempDir::new().unwrap();
        let ndb = super::super::init_ndb(tmp.path()).unwrap();
        let spambox =
            super::super::init_ndb_at_path(&tmp.path().join("nostrdb_spambox"), None).unwrap();

        let root_keys = nostr::Keys::generate();
        let root_pk = root_keys.public_key().to_bytes();
        super::super::set_social_graph_root(&ndb, &root_pk);
        tokio::time::sleep(Duration::from_millis(100)).await;

        let crawler = SocialGraphCrawler::new(Arc::clone(&ndb), root_keys.clone(), vec![], 2)
            .with_spambox(Arc::clone(&spambox));

        let target_keys = nostr::Keys::generate();
        let target_pk = target_keys.public_key().to_bytes();
        let follow_tag = Tag::public_key(PublicKey::from_slice(&target_pk).unwrap());
        let event = EventBuilder::new(Kind::ContactList, "", vec![follow_tag])
            .to_event(&root_keys)
            .unwrap();

        crawler.handle_incoming_event(&event);

        assert!(wait_for_follow(&ndb, &root_pk, &target_pk).await);
        assert!(!wait_for_follow(&spambox, &root_pk, &target_pk).await);
    }

    #[tokio::test]
    async fn test_crawler_no_relays() {
        let tmp = TempDir::new().unwrap();
        let ndb = {
            let _guard = super::super::test_lock();
            super::super::init_ndb(tmp.path()).unwrap()
        };
        let keys = nostr::Keys::generate();
        let crawler = SocialGraphCrawler::new(ndb, keys, vec![], 2);
        let (_tx, rx) = watch::channel(false);
        // Should return immediately with no relays
        crawler.crawl(rx).await;
    }

    #[tokio::test]
    async fn test_crawler_shutdown_signal() {
        let tmp = TempDir::new().unwrap();
        let ndb = {
            let _guard = super::super::test_lock();
            super::super::init_ndb(tmp.path()).unwrap()
        };
        let keys = nostr::Keys::generate();
        let crawler = SocialGraphCrawler::new(ndb, keys, vec!["wss://localhost:1".to_string()], 2);
        let (_tx, rx) = watch::channel(true); // Already shutdown
        crawler.crawl(rx).await;
    }

    #[tokio::test]
    async fn test_collect_missing_root_follows_skips_known_and_fetched() {
        let _guard = super::super::test_lock();
        let tmp = TempDir::new().unwrap();
        let ndb = super::super::init_ndb(tmp.path()).unwrap();

        let root_keys = nostr::Keys::generate();
        let root_pk = root_keys.public_key().to_bytes();
        super::super::set_social_graph_root(&ndb, &root_pk);
        tokio::time::sleep(Duration::from_millis(100)).await;

        let crawler = SocialGraphCrawler::new(Arc::clone(&ndb), root_keys.clone(), vec![], 2);

        let known_keys = nostr::Keys::generate();
        let known_pk = known_keys.public_key().to_bytes();
        let known_follow_tag = Tag::public_key(PublicKey::from_slice(&root_pk).unwrap());
        let known_event = EventBuilder::new(Kind::ContactList, "", vec![known_follow_tag])
            .to_event(&known_keys)
            .unwrap();
        crawler.ingest_event_into(&ndb, "test", &known_event);
        assert!(wait_for_follow(&ndb, &known_pk, &root_pk).await);

        let missing_keys = nostr::Keys::generate();
        let missing_pk = missing_keys.public_key().to_bytes();

        let fetched_keys = nostr::Keys::generate();
        let fetched_pk = fetched_keys.public_key().to_bytes();

        let tags = vec![
            Tag::public_key(PublicKey::from_slice(&known_pk).unwrap()),
            Tag::public_key(PublicKey::from_slice(&missing_pk).unwrap()),
            Tag::public_key(PublicKey::from_slice(&fetched_pk).unwrap()),
        ];
        let root_event = EventBuilder::new(Kind::ContactList, "", tags)
            .to_event(&root_keys)
            .unwrap();

        let mut fetched = HashSet::new();
        fetched.insert(fetched_pk);

        let missing = crawler.collect_missing_root_follows(&root_event, &mut fetched);

        assert_eq!(missing.len(), 1);
        assert_eq!(missing[0], missing_pk);
        assert!(fetched.contains(&known_pk));
        assert!(fetched.contains(&missing_pk));
        assert!(fetched.contains(&fetched_pk));
    }

    #[tokio::test]
    async fn test_collect_missing_root_follows_ignores_non_root() {
        let _guard = super::super::test_lock();
        let tmp = TempDir::new().unwrap();
        let ndb = super::super::init_ndb(tmp.path()).unwrap();

        let root_keys = nostr::Keys::generate();
        let root_pk = root_keys.public_key().to_bytes();
        super::super::set_social_graph_root(&ndb, &root_pk);
        tokio::time::sleep(Duration::from_millis(100)).await;

        let crawler = SocialGraphCrawler::new(Arc::clone(&ndb), root_keys.clone(), vec![], 2);

        let other_keys = nostr::Keys::generate();
        let other_pk = other_keys.public_key().to_bytes();
        let tag = Tag::public_key(PublicKey::from_slice(&other_pk).unwrap());
        let event = EventBuilder::new(Kind::ContactList, "", vec![tag])
            .to_event(&other_keys)
            .unwrap();

        let mut fetched = HashSet::new();
        let missing = crawler.collect_missing_root_follows(&event, &mut fetched);

        assert!(missing.is_empty());
        assert!(fetched.is_empty());
    }
}