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