Skip to main content

hashtree_cli/nostrdb_integration/
mod.rs

1//! Nostrdb integration for social graph-based access control and peer classification.
2
3pub mod access;
4pub mod crawler;
5pub mod snapshot;
6
7pub use nostrdb_social::Ndb;
8use nostrdb_social::{Config as NdbConfig, Transaction};
9use std::collections::HashMap;
10use std::path::Path;
11use std::sync::Arc;
12
13#[cfg(test)]
14use std::sync::{Mutex, MutexGuard, OnceLock};
15
16#[cfg(test)]
17pub type TestLockGuard = MutexGuard<'static, ()>;
18
19#[cfg(test)]
20static NDB_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
21
22#[cfg(test)]
23pub fn test_lock() -> TestLockGuard {
24    NDB_TEST_LOCK.get_or_init(|| Mutex::new(())).lock().unwrap()
25}
26
27pub use access::SocialGraphAccessControl;
28pub use crawler::SocialGraphCrawler;
29
30/// Social graph statistics
31#[derive(Debug, Clone, Default, serde::Serialize)]
32pub struct SocialGraphStats {
33    pub root: Option<String>,
34    pub total_follows: usize,
35    pub max_depth: u32,
36    pub enabled: bool,
37}
38
39/// Initialize nostrdb_social with the given data directory.
40pub fn init_ndb(data_dir: &Path) -> anyhow::Result<Arc<Ndb>> {
41    init_ndb_with_mapsize(data_dir, None)
42}
43
44/// Initialize nostrdb_social with optional mapsize (bytes).
45pub fn init_ndb_with_mapsize(
46    data_dir: &Path,
47    mapsize_bytes: Option<u64>,
48) -> anyhow::Result<Arc<Ndb>> {
49    let ndb_dir = data_dir.join("nostrdb_social");
50    init_ndb_at_path(&ndb_dir, mapsize_bytes)
51}
52
53/// Initialize nostrdb_social at a specific directory (used for spambox).
54pub fn init_ndb_at_path(db_dir: &Path, mapsize_bytes: Option<u64>) -> anyhow::Result<Arc<Ndb>> {
55    std::fs::create_dir_all(db_dir)?;
56    let mut config = NdbConfig::new().set_ingester_threads(2);
57    if let Some(bytes) = mapsize_bytes {
58        let mapsize = usize::try_from(bytes).unwrap_or(usize::MAX);
59        config = config.set_mapsize(mapsize);
60    }
61    let ndb = Ndb::new(db_dir.to_str().unwrap_or("."), &config)?;
62    Ok(Arc::new(ndb))
63}
64
65/// Set the social graph root pubkey.
66pub fn set_social_graph_root(ndb: &Ndb, pk_bytes: &[u8; 32]) {
67    nostrdb_social::socialgraph::set_root(ndb, pk_bytes);
68}
69
70/// Get the follow distance for a pubkey from the social graph root.
71/// Returns None if the pubkey is not in the social graph.
72pub fn get_follow_distance(ndb: &Ndb, pk_bytes: &[u8; 32]) -> Option<u32> {
73    let txn = Transaction::new(ndb).ok()?;
74    let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, pk_bytes);
75    if distance >= 1000 {
76        None
77    } else {
78        Some(distance)
79    }
80}
81
82/// Get the list of pubkeys that a given pubkey follows.
83pub fn get_follows(ndb: &Ndb, pk_bytes: &[u8; 32]) -> Vec<[u8; 32]> {
84    let txn = match Transaction::new(ndb) {
85        Ok(t) => t,
86        Err(_) => return Vec::new(),
87    };
88    nostrdb_social::socialgraph::get_followed(&txn, ndb, pk_bytes, 10000)
89}
90
91fn clamp_socialgraph_list(count: usize) -> usize {
92    let max = i32::MAX as usize;
93    if count > max {
94        max
95    } else {
96        count
97    }
98}
99
100/// Check if a user is overmuted based on muters vs followers at the closest distance
101/// where there is any opinion. Mirrors nostr-social-graph logic.
102pub fn is_overmuted(ndb: &Ndb, root_pk: &[u8; 32], user_pk: &[u8; 32], threshold: f64) -> bool {
103    if threshold <= 0.0 {
104        return false;
105    }
106    if user_pk == root_pk {
107        return false;
108    }
109
110    let txn = match Transaction::new(ndb) {
111        Ok(t) => t,
112        Err(_) => return false,
113    };
114
115    let muter_count = nostrdb_social::socialgraph::muter_count(&txn, ndb, user_pk);
116    if muter_count == 0 {
117        return false;
118    }
119
120    if nostrdb_social::socialgraph::is_muting(&txn, ndb, root_pk, user_pk) {
121        return true;
122    }
123
124    let follower_count = nostrdb_social::socialgraph::follower_count(&txn, ndb, user_pk);
125
126    let mut stats: HashMap<u32, (usize, usize)> = HashMap::new();
127
128    let followers = nostrdb_social::socialgraph::get_followers(
129        &txn,
130        ndb,
131        user_pk,
132        clamp_socialgraph_list(follower_count),
133    );
134    for follower_pk in followers {
135        let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, &follower_pk);
136        if distance >= 1000 {
137            continue;
138        }
139        let entry = stats.entry(distance).or_insert((0, 0));
140        entry.0 += 1;
141    }
142
143    let muters = nostrdb_social::socialgraph::get_muters(
144        &txn,
145        ndb,
146        user_pk,
147        clamp_socialgraph_list(muter_count),
148    );
149    for muter_pk in muters {
150        let distance = nostrdb_social::socialgraph::get_follow_distance(&txn, ndb, &muter_pk);
151        if distance >= 1000 {
152            continue;
153        }
154        let entry = stats.entry(distance).or_insert((0, 0));
155        entry.1 += 1;
156    }
157
158    let mut distances: Vec<u32> = stats.keys().cloned().collect();
159    distances.sort_unstable();
160
161    for distance in distances {
162        let (followers, muters) = stats[&distance];
163        if followers + muters > 0 {
164            return (muters as f64) * threshold > (followers as f64);
165        }
166    }
167
168    false
169}
170
171/// Ingest a Nostr event JSON string into nostrdb_social.
172/// Wraps the event in relay format: ["EVENT","sub_id",{...}]
173pub fn ingest_event(ndb: &Ndb, sub_id: &str, event_json: &str) {
174    let relay_msg = format!(r#"["EVENT","{}",{}]"#, sub_id, event_json);
175    if let Err(e) = ndb.process_event(&relay_msg) {
176        tracing::warn!("Failed to ingest event into nostrdb_social: {}", e);
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183    use tempfile::TempDir;
184
185    #[test]
186    fn test_init_ndb() {
187        let _guard = test_lock();
188        let tmp = TempDir::new().unwrap();
189        let ndb = init_ndb(tmp.path()).unwrap();
190        assert!(Arc::strong_count(&ndb) == 1);
191    }
192
193    #[test]
194    fn test_set_root_and_get_follow_distance() {
195        let _guard = test_lock();
196        let tmp = TempDir::new().unwrap();
197        let ndb = init_ndb(tmp.path()).unwrap();
198        let root_pk = [1u8; 32];
199        set_social_graph_root(&ndb, &root_pk);
200        // Give nostrdb_social a moment to process the root setting
201        std::thread::sleep(std::time::Duration::from_millis(100));
202        let dist = get_follow_distance(&ndb, &root_pk);
203        assert_eq!(dist, Some(0));
204    }
205
206    #[test]
207    fn test_unknown_pubkey_follow_distance() {
208        let _guard = test_lock();
209        let tmp = TempDir::new().unwrap();
210        let ndb = init_ndb(tmp.path()).unwrap();
211        let root_pk = [1u8; 32];
212        set_social_graph_root(&ndb, &root_pk);
213        std::thread::sleep(std::time::Duration::from_millis(100));
214        let unknown_pk = [2u8; 32];
215        assert_eq!(get_follow_distance(&ndb, &unknown_pk), None);
216    }
217
218    #[test]
219    fn test_ingest_event_no_panic() {
220        let _guard = test_lock();
221        let tmp = TempDir::new().unwrap();
222        let ndb = init_ndb(tmp.path()).unwrap();
223        // Pass invalid event - should not panic, just log warning
224        ingest_event(&ndb, "sub1", r#"{"kind":1,"content":"hello"}"#);
225    }
226
227    #[test]
228    fn test_get_follows_empty() {
229        let _guard = test_lock();
230        let tmp = TempDir::new().unwrap();
231        let ndb = init_ndb(tmp.path()).unwrap();
232        let pk = [1u8; 32];
233        assert!(get_follows(&ndb, &pk).is_empty());
234    }
235}