Something went wrong. Try again.
fork of hydrant
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211use fjall::{Keyspace, OwnedWriteBatch};use jacquard_common::types::string::Did;use miette::{IntoDiagnostic, Result};
use crate::db::keys::SEP;use crate::db::types::TrimmedDid;use crate::filter::{FilterConfig, FilterMode};use crate::patch::SetUpdate;
pub const MODE_KEY: &[u8] = b"m";pub const SIGNAL_PREFIX: u8 = b's';pub const COLLECTION_PREFIX: u8 = b'c';pub const EXCLUDE_PREFIX: u8 = b'x';
pub fn signal_key(val: &str) -> Result<Vec<u8>> { let mut key = Vec::with_capacity(2 + val.len()); key.push(SIGNAL_PREFIX); key.push(SEP); key.extend_from_slice(val.as_bytes()); Ok(key)}
pub fn collection_key(val: &str) -> Result<Vec<u8>> { let mut key = Vec::with_capacity(2 + val.len()); key.push(COLLECTION_PREFIX); key.push(SEP); key.extend_from_slice(val.as_bytes()); Ok(key)}
pub fn exclude_key(val: &str) -> Result<Vec<u8>> { let did = Did::<&str>::new(val).into_diagnostic()?; let trimmed = TrimmedDid::from_did_str(did.as_str()); let mut key = Vec::with_capacity(2 + trimmed.len()); key.push(EXCLUDE_PREFIX); key.push(SEP); trimmed.write_to_vec(&mut key); Ok(key)}
pub fn apply_patch( batch: &mut OwnedWriteBatch, ks: &Keyspace, mode: Option<FilterMode>, signals: Option<SetUpdate>, collections: Option<SetUpdate>, excludes: Option<SetUpdate>,) -> Result<()> { if let Some(mode) = mode { batch.insert(ks, MODE_KEY, rmp_serde::to_vec(&mode).into_diagnostic()?); }
apply_set_update(batch, ks, SIGNAL_PREFIX, signals)?; apply_set_update(batch, ks, COLLECTION_PREFIX, collections)?; apply_set_update(batch, ks, EXCLUDE_PREFIX, excludes)?;
Ok(())}
fn apply_set_update( batch: &mut OwnedWriteBatch, ks: &Keyspace, prefix: u8, update: Option<SetUpdate>,) -> Result<()> { let Some(update) = update else { return Ok(()) };
let key_fn = match prefix { SIGNAL_PREFIX => signal_key, COLLECTION_PREFIX => collection_key, EXCLUDE_PREFIX => exclude_key, _ => unreachable!(), };
match update { SetUpdate::Set(values) => { let scan_prefix = [prefix, SEP]; for guard in ks.prefix(scan_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; batch.remove(ks, k); } for val in values { batch.insert(ks, key_fn(&val)?, []); } } SetUpdate::Patch(map) => { for (val, add) in map { let key = key_fn(&val)?; if add { batch.insert(ks, key, []); } else { batch.remove(ks, key); } } } }
Ok(())}
pub fn load(ks: &Keyspace) -> Result<FilterConfig> { let mode = ks .get(MODE_KEY) .into_diagnostic()? .map(|v| rmp_serde::from_slice(&v).into_diagnostic()) .transpose()? .unwrap_or(FilterMode::Filter);
let mut config = FilterConfig::new(mode);
let signal_prefix = [SIGNAL_PREFIX, SEP]; for guard in ks.prefix(signal_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; let val = std::str::from_utf8(&k[signal_prefix.len()..]).into_diagnostic()?; config.signals.push(val.into()); }
let col_prefix = [COLLECTION_PREFIX, SEP]; for guard in ks.prefix(col_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; let val = std::str::from_utf8(&k[col_prefix.len()..]).into_diagnostic()?; config.collections.push(val.into()); }
Ok(config)}
pub fn read_set(ks: &Keyspace, prefix: u8) -> Result<Vec<String>> { let scan_prefix = [prefix, SEP]; let mut out = Vec::new(); for guard in ks.prefix(scan_prefix) { let (k, _) = guard.into_inner().into_diagnostic()?; let val_bytes = &k[2..]; let val = if prefix == EXCLUDE_PREFIX { TrimmedDid::try_from(val_bytes)?.to_did().to_string() } else { std::str::from_utf8(val_bytes).into_diagnostic()?.to_owned() }; out.push(val); } Ok(out)}
#[cfg(test)]mod tests { use smol_str::SmolStr;
use super::*;
#[test] fn test_filter_keys() { assert_eq!( signal_key("app.bsky.feed.like").unwrap(), b"s|app.bsky.feed.like" ); assert_eq!( collection_key("app.bsky.feed.post").unwrap(), b"c|app.bsky.feed.post" ); }
#[test] fn test_exclude_key_trimmed() { let did = "did:plc:yk4q3id7id6p5z3bypvshc64"; let key = exclude_key(did).unwrap(); assert_eq!(key[0], EXCLUDE_PREFIX); assert_eq!(key[1], SEP); // TAG_PLC (1) + 15 bytes assert_eq!(key.len(), 2 + 1 + 15);
let parsed = TrimmedDid::try_from(&key[2..]).unwrap(); assert_eq!(parsed.to_did().as_str(), did); }
#[test] fn test_apply_and_load() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let keyspace = fjall::Database::builder(tmp.path()) .open() .into_diagnostic()?; let ks = keyspace .keyspace("filter", Default::default) .into_diagnostic()?;
let mut batch = keyspace.batch(); let signals = SetUpdate::Set(vec!["a.b.c".to_string()]); let collections = SetUpdate::Set(vec!["d.e.f".to_string()]); let excludes = SetUpdate::Set(vec!["did:plc:yk4q3id7id6p5z3bypvshc64".to_string()]);
apply_patch( &mut batch, &ks, Some(FilterMode::Filter), Some(signals), Some(collections), Some(excludes), )?; batch.commit().into_diagnostic()?;
let config = load(&ks)?; assert_eq!(config.mode, FilterMode::Filter); assert_eq!(config.signals, vec![SmolStr::new("a.b.c")]); assert_eq!(config.collections, vec![SmolStr::new("d.e.f")]);
let excludes = read_set(&ks, EXCLUDE_PREFIX)?; assert_eq!(excludes, vec!["did:plc:yk4q3id7id6p5z3bypvshc64"]);
Ok(()) }}