Skip to main content

hashtree_cli/socialgraph/
crawler.rs

1use futures::stream::{self, StreamExt};
2use std::collections::HashSet;
3use std::path::PathBuf;
4use std::sync::Arc;
5use std::time::Duration;
6
7use nostr::Timestamp;
8use tokio::sync::watch;
9use tokio::task::JoinHandle;
10
11use super::{
12    read_local_list_file_state, sync_local_list_files_force, sync_local_list_files_if_changed,
13    LocalListFileState, SocialGraphBackend,
14};
15
16const DEFAULT_AUTHOR_BATCH_SIZE: usize = 500;
17const DEFAULT_CONCURRENT_BATCHES: usize = 4;
18const GRAPH_FETCH_TIMEOUT: Duration = Duration::from_secs(15);
19const RELAY_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
20pub(crate) const SOCIALGRAPH_RELAY_EVENT_MAX_SIZE: u32 = 512 * 1024;
21#[cfg(not(test))]
22const CRAWLER_STARTUP_DELAY: Duration = Duration::from_secs(5);
23#[cfg(test)]
24const CRAWLER_STARTUP_DELAY: Duration = Duration::from_millis(50);
25#[cfg(not(test))]
26const LOCAL_LIST_POLL_INTERVAL: Duration = Duration::from_secs(5);
27#[cfg(test)]
28const LOCAL_LIST_POLL_INTERVAL: Duration = Duration::from_millis(100);
29
30pub struct SocialGraphTaskHandles {
31    pub shutdown_tx: watch::Sender<bool>,
32    pub crawl_handle: JoinHandle<()>,
33    pub local_list_handle: JoinHandle<()>,
34}
35
36pub struct SocialGraphCrawler {
37    graph_store: Arc<dyn SocialGraphBackend>,
38    spambox: Option<Arc<dyn SocialGraphBackend>>,
39    keys: nostr::Keys,
40    relays: Vec<String>,
41    max_depth: u32,
42    author_batch_size: usize,
43    concurrent_batches: usize,
44    full_recrawl: bool,
45    known_since: Option<Timestamp>,
46}
47
48impl SocialGraphCrawler {
49    pub fn new(
50        graph_store: Arc<dyn SocialGraphBackend>,
51        keys: nostr::Keys,
52        relays: Vec<String>,
53        max_depth: u32,
54    ) -> Self {
55        Self {
56            graph_store,
57            spambox: None,
58            keys,
59            relays,
60            max_depth,
61            author_batch_size: DEFAULT_AUTHOR_BATCH_SIZE,
62            concurrent_batches: DEFAULT_CONCURRENT_BATCHES,
63            full_recrawl: false,
64            known_since: None,
65        }
66    }
67
68    pub fn with_spambox(mut self, spambox: Arc<dyn SocialGraphBackend>) -> Self {
69        self.spambox = Some(spambox);
70        self
71    }
72
73    pub fn with_author_batch_size(mut self, author_batch_size: usize) -> Self {
74        self.author_batch_size = author_batch_size.max(1);
75        self
76    }
77
78    pub fn with_concurrent_batches(mut self, concurrent_batches: usize) -> Self {
79        self.concurrent_batches = concurrent_batches.max(1);
80        self
81    }
82
83    pub fn with_full_recrawl(mut self, full_recrawl: bool) -> Self {
84        self.full_recrawl = full_recrawl;
85        self
86    }
87
88    pub fn with_known_since(mut self, known_since: Option<u64>) -> Self {
89        self.known_since = known_since.map(Timestamp::from);
90        self
91    }
92
93    fn is_within_social_graph(&self, pk_bytes: &[u8; 32]) -> bool {
94        if pk_bytes == &self.keys.public_key().to_bytes() {
95            return true;
96        }
97
98        super::get_follow_distance(self.graph_store.as_ref(), pk_bytes)
99            .map(|distance| distance <= self.max_depth)
100            .unwrap_or(false)
101    }
102
103    fn ingest_events_into(
104        &self,
105        graph_store: &(impl SocialGraphBackend + ?Sized),
106        events: &[nostr::Event],
107    ) {
108        if let Err(err) = super::ingest_parsed_events(graph_store, events) {
109            tracing::debug!("Failed to ingest crawler event: {}", err);
110        }
111    }
112
113    #[allow(deprecated)]
114    fn collect_missing_root_follows(
115        &self,
116        event: &nostr::Event,
117        fetched_contact_lists: &mut HashSet<[u8; 32]>,
118    ) -> Vec<[u8; 32]> {
119        if self.max_depth < 2 || event.kind != nostr::Kind::ContactList {
120            return Vec::new();
121        }
122
123        let root_pk = self.keys.public_key().to_bytes();
124        if event.pubkey.to_bytes() != root_pk {
125            return Vec::new();
126        }
127
128        let mut missing = Vec::new();
129        for tag in event.tags.iter() {
130            if let Some(nostr::TagStandard::PublicKey { public_key, .. }) = tag.as_standardized() {
131                let pk_bytes = public_key.to_bytes();
132                if fetched_contact_lists.contains(&pk_bytes) {
133                    continue;
134                }
135
136                let existing_follows = super::get_follows(self.graph_store.as_ref(), &pk_bytes);
137                if !existing_follows.is_empty() {
138                    fetched_contact_lists.insert(pk_bytes);
139                    continue;
140                }
141
142                fetched_contact_lists.insert(pk_bytes);
143                missing.push(pk_bytes);
144            }
145        }
146
147        missing
148    }
149
150    fn graph_filter_for_pubkeys(
151        pubkeys: &[[u8; 32]],
152        since: Option<Timestamp>,
153    ) -> Option<nostr::Filter> {
154        let authors = pubkeys
155            .iter()
156            .filter_map(|pk_bytes| nostr::PublicKey::from_slice(pk_bytes).ok())
157            .collect::<Vec<_>>();
158        if authors.is_empty() {
159            return None;
160        }
161
162        let mut filter = nostr::Filter::new()
163            .authors(authors)
164            .kinds(vec![nostr::Kind::ContactList, nostr::Kind::MuteList]);
165        if let Some(since) = since {
166            filter = filter.since(since);
167        }
168
169        Some(filter)
170    }
171
172    async fn fetch_graph_events_for_pubkeys(
173        &self,
174        client: &nostr_sdk::Client,
175        pubkeys: &[[u8; 32]],
176        since: Option<Timestamp>,
177    ) -> Vec<nostr::Event> {
178        let Some(filter) = Self::graph_filter_for_pubkeys(pubkeys, since) else {
179            return Vec::new();
180        };
181
182        match client.fetch_events(filter, GRAPH_FETCH_TIMEOUT).await {
183            Ok(events) => events.to_vec(),
184            Err(err) => {
185                tracing::debug!(
186                    "Failed to fetch graph events for {} authors: {}",
187                    pubkeys.len(),
188                    err
189                );
190                Vec::new()
191            }
192        }
193    }
194
195    async fn fetch_contact_lists_for_pubkeys(
196        &self,
197        client: &nostr_sdk::Client,
198        pubkeys: &[[u8; 32]],
199        since: Option<Timestamp>,
200        shutdown_rx: &watch::Receiver<bool>,
201    ) {
202        let chunk_futures = stream::iter(
203            pubkeys
204                .chunks(self.author_batch_size)
205                .map(|chunk| chunk.to_vec())
206                .collect::<Vec<_>>(),
207        )
208        .map(|chunk| async move {
209            self.fetch_graph_events_for_pubkeys(client, &chunk, since)
210                .await
211        });
212
213        let mut in_flight = chunk_futures.buffer_unordered(self.concurrent_batches);
214        while let Some(events) = in_flight.next().await {
215            if *shutdown_rx.borrow() {
216                break;
217            }
218            if let Err(err) = super::ingest_graph_parsed_events(self.graph_store.as_ref(), &events)
219            {
220                tracing::debug!("Failed to ingest crawler graph batch: {}", err);
221            }
222        }
223    }
224
225    fn authors_to_fetch_at_distance(&self, distance: u32) -> (Vec<[u8; 32]>, Vec<[u8; 32]>) {
226        let Ok(users) = self.graph_store.users_by_follow_distance(distance) else {
227            return (Vec::new(), Vec::new());
228        };
229        if self.full_recrawl {
230            return (users, Vec::new());
231        }
232
233        let mut new_authors = Vec::new();
234        let mut known_authors = Vec::new();
235        let refresh_known_authors = self.known_since.is_some();
236        for pk_bytes in users {
237            match self
238                .graph_store
239                .follow_list_created_at(&pk_bytes)
240                .ok()
241                .flatten()
242            {
243                Some(_) if refresh_known_authors => known_authors.push(pk_bytes),
244                Some(_) => {}
245                None => new_authors.push(pk_bytes),
246            }
247        }
248        (new_authors, known_authors)
249    }
250
251    async fn connect_client(&self) -> Option<nostr_sdk::Client> {
252        use nostr::nips::nip19::ToBech32;
253
254        let secret = self.keys.secret_key().to_bech32().unwrap_or_default();
255        let Ok(sdk_keys) = nostr_sdk::Keys::parse(&secret) else {
256            return None;
257        };
258
259        let mut relay_limits = nostr_sdk::pool::RelayLimits::default();
260        relay_limits.events.max_size = Some(SOCIALGRAPH_RELAY_EVENT_MAX_SIZE);
261        let client = nostr_sdk::Client::builder()
262            .signer(sdk_keys)
263            .opts(nostr_sdk::ClientOptions::new().relay_limits(relay_limits))
264            .build();
265        for relay in &self.relays {
266            if let Err(err) = client.add_relay(relay).await {
267                tracing::warn!("Failed to add relay {}: {}", relay, err);
268            }
269        }
270        let _ = tokio::time::timeout(RELAY_CONNECT_TIMEOUT, client.connect()).await;
271        Some(client)
272    }
273
274    async fn sync_graph_once(
275        &self,
276        client: &nostr_sdk::Client,
277        shutdown_rx: &watch::Receiver<bool>,
278        fetched_contact_lists: &mut HashSet<[u8; 32]>,
279    ) {
280        for distance in 0..=self.max_depth {
281            if *shutdown_rx.borrow() {
282                break;
283            }
284
285            let (new_authors, known_authors) = self.authors_to_fetch_at_distance(distance);
286            if new_authors.is_empty() && known_authors.is_empty() {
287                continue;
288            }
289
290            tracing::debug!(
291                "Social graph sync distance={} new_authors={} known_authors={}",
292                distance,
293                new_authors.len(),
294                known_authors.len()
295            );
296
297            if !new_authors.is_empty() {
298                for pk_bytes in &new_authors {
299                    fetched_contact_lists.insert(*pk_bytes);
300                }
301                self.fetch_contact_lists_for_pubkeys(client, &new_authors, None, shutdown_rx)
302                    .await;
303            }
304
305            if !known_authors.is_empty() {
306                for pk_bytes in &known_authors {
307                    fetched_contact_lists.insert(*pk_bytes);
308                }
309                self.fetch_contact_lists_for_pubkeys(
310                    client,
311                    &known_authors,
312                    self.known_since,
313                    shutdown_rx,
314                )
315                .await;
316            }
317        }
318    }
319
320    pub async fn warm_once(&self) {
321        if self.relays.is_empty() {
322            tracing::warn!("Social graph crawler: no relays configured, skipping");
323            return;
324        }
325
326        let Some(client) = self.connect_client().await else {
327            return;
328        };
329        let (_shutdown_tx, shutdown_rx) = watch::channel(false);
330        let mut fetched_contact_lists: HashSet<[u8; 32]> = HashSet::new();
331        self.sync_graph_once(&client, &shutdown_rx, &mut fetched_contact_lists)
332            .await;
333        let _ = client.disconnect().await;
334    }
335
336    pub(crate) fn handle_incoming_event(&self, event: &nostr::Event) {
337        let is_contact_list = event.kind == nostr::Kind::ContactList;
338        let is_mute_list = event.kind == nostr::Kind::MuteList;
339        if !is_contact_list && !is_mute_list {
340            return;
341        }
342
343        let pk_bytes = event.pubkey.to_bytes();
344        if self.is_within_social_graph(&pk_bytes) {
345            self.ingest_events_into(self.graph_store.as_ref(), std::slice::from_ref(event));
346            return;
347        }
348
349        if let Some(spambox) = &self.spambox {
350            self.ingest_events_into(spambox.as_ref(), std::slice::from_ref(event));
351        }
352    }
353
354    #[allow(deprecated)]
355    pub async fn crawl(&self, shutdown_rx: watch::Receiver<bool>) {
356        use nostr_sdk::prelude::RelayPoolNotification;
357
358        if self.relays.is_empty() {
359            tracing::warn!("Social graph crawler: no relays configured, skipping");
360            return;
361        }
362
363        let mut shutdown_rx = shutdown_rx;
364        if *shutdown_rx.borrow() {
365            return;
366        }
367
368        let Some(client) = self.connect_client().await else {
369            return;
370        };
371
372        let mut fetched_contact_lists: HashSet<[u8; 32]> = HashSet::new();
373        self.sync_graph_once(&client, &shutdown_rx, &mut fetched_contact_lists)
374            .await;
375
376        let filter = nostr::Filter::new()
377            .kinds(vec![nostr::Kind::ContactList, nostr::Kind::MuteList])
378            .since(nostr::Timestamp::now());
379
380        let _ = client.subscribe(filter, None).await;
381
382        let mut notifications = client.notifications();
383        loop {
384            tokio::select! {
385                _ = shutdown_rx.changed() => {
386                    if *shutdown_rx.borrow() {
387                        break;
388                    }
389                }
390                notification = notifications.recv() => {
391                    match notification {
392                        Ok(RelayPoolNotification::Event { event, .. }) => {
393                            self.handle_incoming_event(&event);
394                            let missing = self.collect_missing_root_follows(&event, &mut fetched_contact_lists);
395                            if !missing.is_empty() {
396                                self.fetch_contact_lists_for_pubkeys(&client, &missing, None, &shutdown_rx).await;
397                            }
398                        }
399                        Ok(_) => {}
400                        Err(err) => {
401                            tracing::warn!("Social graph crawler notification error: {}", err);
402                            break;
403                        }
404                    }
405                }
406            }
407        }
408
409        let _ = client.disconnect().await;
410    }
411}
412
413pub fn spawn_social_graph_tasks(
414    graph_store: Arc<dyn SocialGraphBackend>,
415    keys: nostr::Keys,
416    relays: Vec<String>,
417    max_depth: u32,
418    spambox: Option<Arc<dyn SocialGraphBackend>>,
419    data_dir: PathBuf,
420) -> SocialGraphTaskHandles {
421    let (shutdown_tx, crawl_shutdown_rx) = watch::channel(false);
422    let local_list_shutdown_rx = crawl_shutdown_rx.clone();
423    let crawl_data_dir = data_dir.clone();
424    let local_list_data_dir = data_dir;
425
426    let crawl_handle = tokio::spawn({
427        let graph_store = Arc::clone(&graph_store);
428        let keys = keys.clone();
429        let relays = relays.clone();
430        let spambox = spambox.clone();
431        async move {
432            tokio::time::sleep(CRAWLER_STARTUP_DELAY).await;
433            let mut crawler = SocialGraphCrawler::new(graph_store.clone(), keys, relays, max_depth);
434            if let Some(spambox) = spambox {
435                crawler = crawler.with_spambox(spambox);
436            }
437            if let Err(err) =
438                sync_local_list_files_force(graph_store.as_ref(), &crawl_data_dir, &crawler.keys)
439            {
440                tracing::warn!(
441                    "Failed to sync local social graph lists at startup: {}",
442                    err
443                );
444            }
445            crawler.crawl(crawl_shutdown_rx).await;
446        }
447    });
448
449    let local_list_handle = tokio::spawn({
450        let graph_store = Arc::clone(&graph_store);
451        let keys = keys.clone();
452        let relays = relays.clone();
453        let spambox = spambox.clone();
454        async move {
455            let mut shutdown_rx = local_list_shutdown_rx;
456            let mut state =
457                read_local_list_file_state(&local_list_data_dir).unwrap_or_else(|err| {
458                    tracing::warn!("Failed to read local social graph list state: {}", err);
459                    LocalListFileState::default()
460                });
461            let mut interval = tokio::time::interval(LOCAL_LIST_POLL_INTERVAL);
462            loop {
463                tokio::select! {
464                    _ = shutdown_rx.changed() => {
465                        if *shutdown_rx.borrow() {
466                            break;
467                        }
468                    }
469                    _ = interval.tick() => {
470                        let outcome = match sync_local_list_files_if_changed(
471                            graph_store.as_ref(),
472                            &local_list_data_dir,
473                            &keys,
474                            &mut state,
475                        ) {
476                            Ok(outcome) => outcome,
477                            Err(err) => {
478                                tracing::warn!("Failed to sync local social graph list files: {}", err);
479                                continue;
480                            }
481                        };
482
483                        if outcome.contacts_changed {
484                            let mut crawler =
485                                SocialGraphCrawler::new(graph_store.clone(), keys.clone(), relays.clone(), max_depth);
486                            if let Some(spambox) = spambox.clone() {
487                                crawler = crawler.with_spambox(spambox);
488                            }
489                            crawler.warm_once().await;
490                        }
491                    }
492                }
493            }
494        }
495    });
496
497    SocialGraphTaskHandles {
498        shutdown_tx,
499        crawl_handle,
500        local_list_handle,
501    }
502}
503
504#[cfg(test)]
505mod tests {
506    use super::*;
507    use std::net::TcpListener;
508    use std::sync::Mutex;
509    use std::time::Instant;
510
511    use futures::{SinkExt, StreamExt};
512    use nostr::{EventBuilder, JsonUtil, Kind, PublicKey, Tag};
513    use tempfile::TempDir;
514    use tokio::net::TcpStream;
515    use tokio::sync::broadcast;
516    use tokio_tungstenite::{accept_async, tungstenite::Message};
517
518    macro_rules! event_builder {
519        ($kind:expr, $content:expr $(,)?) => {
520            EventBuilder::new($kind, $content)
521        };
522        ($kind:expr, $content:expr, $tags:expr $(,)?) => {
523            EventBuilder::new($kind, $content).tags($tags)
524        };
525    }
526
527    #[derive(Debug, Default)]
528    struct RelayState {
529        events: Vec<nostr::Event>,
530        request_author_counts: Vec<usize>,
531        requested_authors: Vec<Vec<String>>,
532    }
533
534    struct TestRelay {
535        port: u16,
536        shutdown: broadcast::Sender<()>,
537        state: Arc<Mutex<RelayState>>,
538    }
539
540    impl TestRelay {
541        fn new(events: Vec<nostr::Event>) -> Self {
542            let state = Arc::new(Mutex::new(RelayState {
543                events,
544                request_author_counts: Vec::new(),
545                requested_authors: Vec::new(),
546            }));
547            let (shutdown, _) = broadcast::channel(1);
548
549            let std_listener = TcpListener::bind("127.0.0.1:0").expect("bind relay listener");
550            let port = std_listener.local_addr().expect("local addr").port();
551            std_listener
552                .set_nonblocking(true)
553                .expect("set listener nonblocking");
554
555            let state_for_thread = Arc::clone(&state);
556            let shutdown_for_thread = shutdown.clone();
557            std::thread::spawn(move || {
558                let runtime = tokio::runtime::Builder::new_multi_thread()
559                    .worker_threads(2)
560                    .enable_all()
561                    .build()
562                    .expect("build tokio runtime");
563                runtime.block_on(async move {
564                    let listener = tokio::net::TcpListener::from_std(std_listener)
565                        .expect("tokio listener from std");
566                    let mut shutdown_rx = shutdown_for_thread.subscribe();
567
568                    loop {
569                        tokio::select! {
570                            _ = shutdown_rx.recv() => break,
571                            accept = listener.accept() => {
572                                if let Ok((stream, _)) = accept {
573                                    let state = Arc::clone(&state_for_thread);
574                                    tokio::spawn(async move {
575                                        handle_connection(stream, state).await;
576                                    });
577                                }
578                            }
579                        }
580                    }
581                });
582            });
583
584            std::thread::sleep(Duration::from_millis(100));
585
586            Self {
587                port,
588                shutdown,
589                state,
590            }
591        }
592
593        fn url(&self) -> String {
594            format!("ws://127.0.0.1:{}", self.port)
595        }
596
597        fn request_author_counts(&self) -> Vec<usize> {
598            self.state
599                .lock()
600                .expect("relay state lock")
601                .request_author_counts
602                .clone()
603        }
604
605        fn requested_authors(&self) -> Vec<Vec<String>> {
606            self.state
607                .lock()
608                .expect("relay state lock")
609                .requested_authors
610                .clone()
611        }
612    }
613
614    impl Drop for TestRelay {
615        fn drop(&mut self) {
616            let _ = self.shutdown.send(());
617            std::thread::sleep(Duration::from_millis(50));
618        }
619    }
620
621    fn matching_events(
622        state: &Arc<Mutex<RelayState>>,
623        filters: &[nostr::Filter],
624    ) -> Vec<nostr::Event> {
625        let guard = state.lock().expect("relay state lock");
626        guard
627            .events
628            .iter()
629            .filter(|event| {
630                filters.is_empty()
631                    || filters
632                        .iter()
633                        .any(|filter| filter.match_event(event, Default::default()))
634            })
635            .cloned()
636            .collect()
637    }
638
639    async fn send_relay_message(
640        write: &mut futures::stream::SplitSink<
641            tokio_tungstenite::WebSocketStream<TcpStream>,
642            Message,
643        >,
644        message: nostr::RelayMessage<'_>,
645    ) {
646        let _ = write.send(Message::Text(message.as_json())).await;
647    }
648
649    async fn handle_connection(stream: TcpStream, state: Arc<Mutex<RelayState>>) {
650        let ws_stream = match accept_async(stream).await {
651            Ok(ws) => ws,
652            Err(_) => return,
653        };
654        let (mut write, mut read) = ws_stream.split();
655
656        while let Some(message) = read.next().await {
657            let text = match message {
658                Ok(Message::Text(text)) => text,
659                Ok(Message::Ping(data)) => {
660                    let _ = write.send(Message::Pong(data)).await;
661                    continue;
662                }
663                Ok(Message::Close(_)) => break,
664                _ => continue,
665            };
666
667            let parsed = match nostr::ClientMessage::from_json(text.as_bytes()) {
668                Ok(message) => message,
669                Err(_) => continue,
670            };
671
672            match parsed {
673                nostr::ClientMessage::Req {
674                    subscription_id,
675                    filters,
676                } => {
677                    let subscription_id = subscription_id.into_owned();
678                    let filters = filters
679                        .into_iter()
680                        .map(|filter| filter.into_owned())
681                        .collect::<Vec<_>>();
682                    let author_count = filters
683                        .iter()
684                        .filter_map(|filter| filter.authors.as_ref())
685                        .map(|authors| authors.len())
686                        .sum();
687                    let mut requested_authors = filters
688                        .iter()
689                        .filter_map(|filter| filter.authors.as_ref())
690                        .flat_map(|authors| authors.iter().map(|author| author.to_hex()))
691                        .collect::<Vec<_>>();
692                    requested_authors.sort();
693                    requested_authors.dedup();
694                    {
695                        let mut guard = state.lock().expect("relay state lock");
696                        guard.request_author_counts.push(author_count);
697                        guard.requested_authors.push(requested_authors);
698                    }
699
700                    for event in matching_events(&state, &filters) {
701                        send_relay_message(
702                            &mut write,
703                            nostr::RelayMessage::event(subscription_id.clone(), event),
704                        )
705                        .await;
706                    }
707                    send_relay_message(&mut write, nostr::RelayMessage::eose(subscription_id))
708                        .await;
709                }
710                nostr::ClientMessage::Close(subscription_id) => {
711                    send_relay_message(
712                        &mut write,
713                        nostr::RelayMessage::closed(subscription_id.into_owned(), ""),
714                    )
715                    .await;
716                }
717                _ => {}
718            }
719        }
720    }
721
722    async fn wait_until<F>(timeout: Duration, mut condition: F)
723    where
724        F: FnMut() -> bool,
725    {
726        let start = Instant::now();
727        while start.elapsed() < timeout {
728            if condition() {
729                return;
730            }
731            tokio::time::sleep(Duration::from_millis(50)).await;
732        }
733        panic!("condition not met within {:?}", timeout);
734    }
735
736    fn write_contacts_file(path: &std::path::Path, pubkeys: &[String]) {
737        std::fs::write(
738            path,
739            serde_json::to_string(pubkeys).expect("serialize contacts"),
740        )
741        .expect("write contacts file");
742    }
743
744    #[tokio::test]
745    async fn test_crawler_routes_untrusted_to_spambox() {
746        let _guard = crate::socialgraph::test_lock();
747        let tmp = TempDir::new().unwrap();
748        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
749        let spambox = crate::socialgraph::open_test_social_graph_store_at_path(
750            &tmp.path().join("spambox"),
751            None,
752        )
753        .unwrap();
754
755        let root_keys = nostr::Keys::generate();
756        let root_pk = root_keys.public_key().to_bytes();
757        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
758        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
759        let spambox_backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = spambox.clone();
760
761        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![], 2)
762            .with_spambox(spambox_backend);
763
764        let unknown_keys = nostr::Keys::generate();
765        let follow_tag = Tag::public_key(PublicKey::from_slice(&root_pk).unwrap());
766        let event = event_builder!(Kind::ContactList, "")
767            .tags(vec![follow_tag])
768            .sign_with_keys(&unknown_keys)
769            .unwrap();
770
771        crawler.handle_incoming_event(&event);
772
773        let unknown_pk = unknown_keys.public_key().to_bytes();
774        assert!(crate::socialgraph::get_follows(&graph_store, &unknown_pk).is_empty());
775        assert_eq!(
776            crate::socialgraph::get_follows(&spambox, &unknown_pk),
777            vec![root_pk]
778        );
779    }
780
781    #[tokio::test]
782    #[allow(clippy::await_holding_lock)]
783    async fn test_crawler_batches_graph_fetches_by_author_chunk() {
784        let _guard = crate::socialgraph::test_lock();
785        let tmp = TempDir::new().unwrap();
786        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
787
788        let root_keys = nostr::Keys::generate();
789        let root_pk = root_keys.public_key().to_bytes();
790        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
791        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
792
793        let alice_keys = nostr::Keys::generate();
794        let bob_keys = nostr::Keys::generate();
795        let carol_keys = nostr::Keys::generate();
796
797        let root_event = event_builder!(
798            Kind::ContactList,
799            "",
800            vec![
801                Tag::public_key(alice_keys.public_key()),
802                Tag::public_key(bob_keys.public_key()),
803            ],
804        )
805        .custom_created_at(nostr::Timestamp::from(10))
806        .sign_with_keys(&root_keys)
807        .unwrap();
808        let alice_event = event_builder!(
809            Kind::ContactList,
810            "",
811            vec![Tag::public_key(carol_keys.public_key())],
812        )
813        .custom_created_at(nostr::Timestamp::from(11))
814        .sign_with_keys(&alice_keys)
815        .unwrap();
816        let bob_event = event_builder!(Kind::ContactList, "")
817            .custom_created_at(nostr::Timestamp::from(12))
818            .sign_with_keys(&bob_keys)
819            .unwrap();
820
821        let relay = TestRelay::new(vec![root_event, alice_event, bob_event]);
822        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 2)
823            .with_author_batch_size(2);
824
825        let (shutdown_tx, shutdown_rx) = watch::channel(false);
826        let handle = tokio::spawn(async move {
827            crawler.crawl(shutdown_rx).await;
828        });
829
830        let alice_pk = alice_keys.public_key().to_bytes();
831        let bob_pk = bob_keys.public_key().to_bytes();
832        let carol_pk = carol_keys.public_key().to_bytes();
833        wait_until(Duration::from_secs(5), || {
834            let root_follows = crate::socialgraph::get_follows(&graph_store, &root_pk);
835            let alice_follows = crate::socialgraph::get_follows(&graph_store, &alice_pk);
836            root_follows.contains(&alice_pk)
837                && root_follows.contains(&bob_pk)
838                && alice_follows.contains(&carol_pk)
839        })
840        .await;
841
842        let _ = shutdown_tx.send(true);
843        handle.await.unwrap();
844
845        let author_counts = relay.request_author_counts();
846        assert!(
847            author_counts.iter().any(|count| *count >= 2),
848            "expected batched author REQ, got {:?}",
849            author_counts
850        );
851    }
852
853    #[tokio::test]
854    #[allow(clippy::await_holding_lock)]
855    async fn test_crawler_expands_from_existing_graph_without_root_refetch() {
856        let _guard = crate::socialgraph::test_lock();
857        let tmp = TempDir::new().unwrap();
858        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
859
860        let root_keys = nostr::Keys::generate();
861        let root_pk = root_keys.public_key().to_bytes();
862        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
863        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
864
865        let alice_keys = nostr::Keys::generate();
866        let bob_keys = nostr::Keys::generate();
867        let carol_keys = nostr::Keys::generate();
868
869        let root_event = event_builder!(
870            Kind::ContactList,
871            "",
872            vec![
873                Tag::public_key(alice_keys.public_key()),
874                Tag::public_key(bob_keys.public_key()),
875            ],
876        )
877        .custom_created_at(nostr::Timestamp::from(10))
878        .sign_with_keys(&root_keys)
879        .unwrap();
880        crate::socialgraph::ingest_parsed_event(&graph_store, &root_event).unwrap();
881
882        let alice_event = event_builder!(
883            Kind::ContactList,
884            "",
885            vec![Tag::public_key(carol_keys.public_key())],
886        )
887        .custom_created_at(nostr::Timestamp::from(11))
888        .sign_with_keys(&alice_keys)
889        .unwrap();
890        let bob_event = event_builder!(Kind::ContactList, "")
891            .custom_created_at(nostr::Timestamp::from(12))
892            .sign_with_keys(&bob_keys)
893            .unwrap();
894
895        let relay = TestRelay::new(vec![alice_event, bob_event]);
896        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 2)
897            .with_author_batch_size(2);
898
899        let (shutdown_tx, shutdown_rx) = watch::channel(false);
900        let handle = tokio::spawn(async move {
901            crawler.crawl(shutdown_rx).await;
902        });
903
904        let alice_pk = alice_keys.public_key().to_bytes();
905        let carol_pk = carol_keys.public_key().to_bytes();
906        wait_until(Duration::from_secs(5), || {
907            crate::socialgraph::get_follows(&graph_store, &alice_pk).contains(&carol_pk)
908        })
909        .await;
910
911        let _ = shutdown_tx.send(true);
912        handle.await.unwrap();
913
914        let author_counts = relay.request_author_counts();
915        let requested_authors = relay.requested_authors();
916        assert!(
917            requested_authors
918                .iter()
919                .flatten()
920                .all(|author| author != &root_keys.public_key().to_hex()),
921            "expected incremental crawl to skip root refetch, got {:?}",
922            requested_authors
923        );
924        assert!(
925            author_counts.iter().any(|count| *count >= 2),
926            "expected batched distance-1 REQ, got {:?}",
927            author_counts
928        );
929    }
930
931    #[tokio::test]
932    #[allow(clippy::await_holding_lock)]
933    async fn test_crawler_full_recrawl_refetches_root() {
934        let _guard = crate::socialgraph::test_lock();
935        let tmp = TempDir::new().unwrap();
936        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
937
938        let root_keys = nostr::Keys::generate();
939        let root_pk = root_keys.public_key().to_bytes();
940        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
941        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
942
943        let alice_keys = nostr::Keys::generate();
944        let bob_keys = nostr::Keys::generate();
945
946        let root_event = event_builder!(
947            Kind::ContactList,
948            "",
949            vec![
950                Tag::public_key(alice_keys.public_key()),
951                Tag::public_key(bob_keys.public_key()),
952            ],
953        )
954        .custom_created_at(nostr::Timestamp::from(10))
955        .sign_with_keys(&root_keys)
956        .unwrap();
957        crate::socialgraph::ingest_parsed_event(&graph_store, &root_event).unwrap();
958
959        let relay = TestRelay::new(vec![root_event]);
960        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 1)
961            .with_full_recrawl(true);
962
963        let (shutdown_tx, shutdown_rx) = watch::channel(false);
964        let handle = tokio::spawn(async move {
965            crawler.crawl(shutdown_rx).await;
966        });
967
968        wait_until(Duration::from_secs(5), || {
969            relay.request_author_counts().contains(&1)
970        })
971        .await;
972
973        let _ = shutdown_tx.send(true);
974        handle.await.unwrap();
975
976        let requested_authors = relay.requested_authors();
977        assert!(
978            requested_authors
979                .iter()
980                .flatten()
981                .any(|author| author == &root_keys.public_key().to_hex()),
982            "expected full recrawl to refetch root, got {:?}",
983            requested_authors
984        );
985    }
986
987    #[tokio::test]
988    #[allow(clippy::await_holding_lock)]
989    async fn test_crawler_revisits_known_authors_with_since_cursor() {
990        let _guard = crate::socialgraph::test_lock();
991        let tmp = TempDir::new().unwrap();
992        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
993
994        let root_keys = nostr::Keys::generate();
995        let root_pk = root_keys.public_key().to_bytes();
996        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
997        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
998
999        let alice_keys = nostr::Keys::generate();
1000        let carol_keys = nostr::Keys::generate();
1001        let dave_keys = nostr::Keys::generate();
1002
1003        let root_event = event_builder!(
1004            Kind::ContactList,
1005            "",
1006            vec![Tag::public_key(alice_keys.public_key())],
1007        )
1008        .custom_created_at(nostr::Timestamp::from(10))
1009        .sign_with_keys(&root_keys)
1010        .unwrap();
1011        crate::socialgraph::ingest_parsed_event(&graph_store, &root_event).unwrap();
1012
1013        let alice_old_event = event_builder!(
1014            Kind::ContactList,
1015            "",
1016            vec![Tag::public_key(carol_keys.public_key())],
1017        )
1018        .custom_created_at(nostr::Timestamp::from(11))
1019        .sign_with_keys(&alice_keys)
1020        .unwrap();
1021        crate::socialgraph::ingest_parsed_event(&graph_store, &alice_old_event).unwrap();
1022
1023        let alice_new_event = event_builder!(
1024            Kind::ContactList,
1025            "",
1026            vec![
1027                Tag::public_key(carol_keys.public_key()),
1028                Tag::public_key(dave_keys.public_key()),
1029            ],
1030        )
1031        .custom_created_at(nostr::Timestamp::from(20))
1032        .sign_with_keys(&alice_keys)
1033        .unwrap();
1034
1035        let relay = TestRelay::new(vec![alice_new_event]);
1036        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 2)
1037            .with_author_batch_size(2)
1038            .with_known_since(Some(15));
1039
1040        let (shutdown_tx, shutdown_rx) = watch::channel(false);
1041        let handle = tokio::spawn(async move {
1042            crawler.crawl(shutdown_rx).await;
1043        });
1044
1045        let alice_pk = alice_keys.public_key().to_bytes();
1046        let dave_pk = dave_keys.public_key().to_bytes();
1047        wait_until(Duration::from_secs(5), || {
1048            crate::socialgraph::get_follows(&graph_store, &alice_pk).contains(&dave_pk)
1049        })
1050        .await;
1051
1052        let _ = shutdown_tx.send(true);
1053        handle.await.unwrap();
1054
1055        let requested_authors = relay.requested_authors();
1056        assert!(
1057            crate::socialgraph::get_follows(&graph_store, &alice_pk).contains(&dave_pk),
1058            "expected incremental crawl to refresh known author follow list"
1059        );
1060        assert!(
1061            requested_authors
1062                .iter()
1063                .flatten()
1064                .any(|author| author == &alice_keys.public_key().to_hex()),
1065            "expected crawler to refetch known author, got {:?}",
1066            requested_authors
1067        );
1068    }
1069
1070    #[tokio::test]
1071    #[allow(clippy::await_holding_lock)]
1072    async fn test_crawler_warm_once_completes_initial_sync_without_shutdown() {
1073        let _guard = crate::socialgraph::test_lock();
1074        let tmp = TempDir::new().unwrap();
1075        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
1076
1077        let root_keys = nostr::Keys::generate();
1078        let root_pk = root_keys.public_key().to_bytes();
1079        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
1080        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1081
1082        let alice_keys = nostr::Keys::generate();
1083        let bob_keys = nostr::Keys::generate();
1084
1085        let root_event = event_builder!(
1086            Kind::ContactList,
1087            "",
1088            vec![Tag::public_key(alice_keys.public_key())],
1089        )
1090        .custom_created_at(nostr::Timestamp::from(10))
1091        .sign_with_keys(&root_keys)
1092        .unwrap();
1093        let alice_event = event_builder!(
1094            Kind::ContactList,
1095            "",
1096            vec![Tag::public_key(bob_keys.public_key())],
1097        )
1098        .custom_created_at(nostr::Timestamp::from(11))
1099        .sign_with_keys(&alice_keys)
1100        .unwrap();
1101
1102        let relay = TestRelay::new(vec![root_event, alice_event]);
1103        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 2)
1104            .with_author_batch_size(2);
1105
1106        let started = std::time::Instant::now();
1107        crawler.warm_once().await;
1108
1109        let alice_pk = alice_keys.public_key().to_bytes();
1110        let bob_pk = bob_keys.public_key().to_bytes();
1111        assert!(
1112            started.elapsed() < Duration::from_secs(5),
1113            "warm_once should complete finite sync promptly"
1114        );
1115        assert!(
1116            crate::socialgraph::get_follows(&graph_store, &alice_pk).contains(&bob_pk),
1117            "expected warm_once to ingest distance-1 follow list"
1118        );
1119    }
1120
1121    #[tokio::test]
1122    async fn test_crawler_warm_once_accepts_large_contact_list_events() {
1123        let _guard = crate::socialgraph::test_lock();
1124        let tmp = TempDir::new().unwrap();
1125        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
1126
1127        let root_keys = nostr::Keys::generate();
1128        let root_pk = root_keys.public_key().to_bytes();
1129        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
1130        let backend: Arc<dyn crate::socialgraph::SocialGraphBackend> = graph_store.clone();
1131
1132        let followed_keys = (0..1_600)
1133            .map(|_| nostr::Keys::generate())
1134            .collect::<Vec<_>>();
1135        let tags = followed_keys
1136            .iter()
1137            .map(|keys| Tag::public_key(keys.public_key()))
1138            .collect::<Vec<_>>();
1139        let root_event = event_builder!(Kind::ContactList, "")
1140            .tags(tags)
1141            .custom_created_at(nostr::Timestamp::from(10))
1142            .sign_with_keys(&root_keys)
1143            .unwrap();
1144        assert!(
1145            root_event.as_json().len() > 70_000,
1146            "test event should exceed nostr-sdk default size limit"
1147        );
1148
1149        let relay = TestRelay::new(vec![root_event]);
1150        let crawler = SocialGraphCrawler::new(backend, root_keys.clone(), vec![relay.url()], 1)
1151            .with_author_batch_size(1);
1152
1153        crawler.warm_once().await;
1154
1155        let follows = crate::socialgraph::get_follows(&graph_store, &root_pk);
1156        let first_pk = followed_keys
1157            .first()
1158            .expect("first followed key")
1159            .public_key()
1160            .to_bytes();
1161        let last_pk = followed_keys
1162            .last()
1163            .expect("last followed key")
1164            .public_key()
1165            .to_bytes();
1166        assert!(
1167            follows.contains(&first_pk) && follows.contains(&last_pk),
1168            "expected warm_once to ingest oversized contact list event"
1169        );
1170    }
1171
1172    #[tokio::test]
1173    #[allow(clippy::await_holding_lock)]
1174    async fn test_spawned_social_graph_tasks_sync_local_contacts_on_startup() {
1175        let _guard = crate::socialgraph::test_lock();
1176        let tmp = TempDir::new().unwrap();
1177        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
1178
1179        let root_keys = nostr::Keys::generate();
1180        let root_pk = root_keys.public_key().to_bytes();
1181        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
1182
1183        let alice_keys = nostr::Keys::generate();
1184        let bob_keys = nostr::Keys::generate();
1185
1186        write_contacts_file(
1187            &tmp.path().join("contacts.json"),
1188            &[alice_keys.public_key().to_hex()],
1189        );
1190
1191        let alice_event = event_builder!(
1192            Kind::ContactList,
1193            "",
1194            vec![Tag::public_key(bob_keys.public_key())],
1195        )
1196        .custom_created_at(nostr::Timestamp::from(11))
1197        .sign_with_keys(&alice_keys)
1198        .unwrap();
1199
1200        let relay = TestRelay::new(vec![alice_event]);
1201        let tasks = spawn_social_graph_tasks(
1202            graph_store.clone(),
1203            root_keys.clone(),
1204            vec![relay.url()],
1205            2,
1206            None,
1207            tmp.path().to_path_buf(),
1208        );
1209
1210        let alice_pk = alice_keys.public_key().to_bytes();
1211        let bob_pk = bob_keys.public_key().to_bytes();
1212        wait_until(Duration::from_secs(5), || {
1213            let root_follows = crate::socialgraph::get_follows(&graph_store, &root_pk);
1214            let alice_follows = crate::socialgraph::get_follows(&graph_store, &alice_pk);
1215            root_follows.contains(&alice_pk) && alice_follows.contains(&bob_pk)
1216        })
1217        .await;
1218
1219        let _ = tasks.shutdown_tx.send(true);
1220        tasks.crawl_handle.abort();
1221        tasks.local_list_handle.abort();
1222    }
1223
1224    #[tokio::test]
1225    #[allow(clippy::await_holding_lock)]
1226    async fn test_spawned_social_graph_tasks_refresh_when_contacts_change() {
1227        let _guard = crate::socialgraph::test_lock();
1228        let tmp = TempDir::new().unwrap();
1229        let graph_store = crate::socialgraph::open_test_social_graph_store(tmp.path()).unwrap();
1230
1231        let root_keys = nostr::Keys::generate();
1232        let root_pk = root_keys.public_key().to_bytes();
1233        crate::socialgraph::set_social_graph_root(&graph_store, &root_pk);
1234
1235        let alice_keys = nostr::Keys::generate();
1236        let bob_keys = nostr::Keys::generate();
1237
1238        let alice_event = event_builder!(
1239            Kind::ContactList,
1240            "",
1241            vec![Tag::public_key(bob_keys.public_key())],
1242        )
1243        .custom_created_at(nostr::Timestamp::from(11))
1244        .sign_with_keys(&alice_keys)
1245        .unwrap();
1246
1247        let relay = TestRelay::new(vec![alice_event]);
1248        let tasks = spawn_social_graph_tasks(
1249            graph_store.clone(),
1250            root_keys.clone(),
1251            vec![relay.url()],
1252            2,
1253            None,
1254            tmp.path().to_path_buf(),
1255        );
1256
1257        let alice_pk = alice_keys.public_key().to_bytes();
1258        let bob_pk = bob_keys.public_key().to_bytes();
1259        assert!(!crate::socialgraph::get_follows(&graph_store, &root_pk).contains(&alice_pk));
1260
1261        write_contacts_file(
1262            &tmp.path().join("contacts.json"),
1263            &[alice_keys.public_key().to_hex()],
1264        );
1265
1266        wait_until(Duration::from_secs(5), || {
1267            let root_follows = crate::socialgraph::get_follows(&graph_store, &root_pk);
1268            let alice_follows = crate::socialgraph::get_follows(&graph_store, &alice_pk);
1269            root_follows.contains(&alice_pk) && alice_follows.contains(&bob_pk)
1270        })
1271        .await;
1272
1273        let _ = tasks.shutdown_tx.send(true);
1274        tasks.crawl_handle.abort();
1275        tasks.local_list_handle.abort();
1276    }
1277}