sneedstr/src/noose/sled.rs
2024-01-26 14:15:53 -06:00

141 lines
4.1 KiB
Rust

use crate::bussy::{channels, Command, Message, PubSub};
use crate::utils::error::Error;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq)]
pub struct BanInfo {
pub pubkey: String,
pub reason: String,
}
// Db Interface
pub struct SledDb {
db: sled::Db,
banned_pubkeys: sled::Tree,
}
impl SledDb {
pub fn new() -> Self {
let db = sled::open("/tmp/sled_db").unwrap();
let banned_pubkeys = db.open_tree("banned_pubkeys").unwrap();
Self { db, banned_pubkeys }
}
pub async fn start(&mut self, pubsub: Arc<PubSub>) -> Result<(), Error> {
let mut subscriber = pubsub.subscribe(channels::MSG_NOOSE).await;
while let Ok(message) = subscriber.recv().await {
log::info!("[Noose] received message: {:?}", message);
let command = match message.content {
Command::SledReqBanUser(ban_info) => match self.ban_user(ban_info).await {
Ok(status) => Command::SledResSuccess(status),
Err(e) => Command::ServiceError(e),
},
Command::SledReqUnbanUser(pubkey) => match self.unban_user(&pubkey).await {
Ok(status) => Command::SledResSuccess(status),
Err(e) => Command::ServiceError(e),
},
Command::SledReqGetBans => match self.get_bans().await {
Ok(bans) => Command::SledResBans(bans),
Err(e) => Command::ServiceError(e),
},
Command::SledReqBanInfo(pubkey) => match self.get_ban_by_pubkey(&pubkey).await {
Ok(ban_info) => Command::SledResBan(ban_info),
Err(e) => Command::ServiceError(e),
},
_ => Command::Noop,
};
if command != Command::Noop {
let channel = message.source;
let message = Message {
source: channels::MSG_SLED,
content: command,
};
log::info!(
"[Sled] publishing new message: {:?} to channel {}",
message,
channel
);
pubsub.publish(channel, message).await;
}
}
Ok(())
}
fn clear_db(&self) -> Result<(), sled::Error> {
self.db.clear()
}
async fn ban_user(&self, ban_info: Box<BanInfo>) -> Result<bool, Error> {
if let Ok(Some(_)) = self
.banned_pubkeys
.insert(ban_info.pubkey, ban_info.reason.as_bytes())
{
return Ok(true);
}
Ok(false)
}
fn is_banned(&self, pubkey: &String) -> bool {
if let Ok(Some(banned)) = self.banned_pubkeys.get(pubkey) {
return true;
}
false
}
async fn unban_user(&self, pubkey: &String) -> Result<bool, Error> {
if self.is_banned(pubkey) {
self.banned_pubkeys.remove(pubkey).unwrap();
return Ok(true);
}
Ok(false)
}
async fn get_bans(&self) -> Result<Vec<BanInfo>, Error> {
let bans: Vec<BanInfo> = self
.banned_pubkeys
.iter()
.filter_map(|row| {
if let Ok((k, v)) = row {
let ban_info = BanInfo {
pubkey: String::from_utf8(k.to_vec()).unwrap(),
reason: String::from_utf8(v.to_vec()).unwrap(),
};
Some(ban_info)
} else {
None
}
})
.collect();
Ok(bans)
}
async fn get_ban_by_pubkey(&self, pubkey: &String) -> Result<Option<BanInfo>, Error> {
if self.is_banned(pubkey) {
if let Ok(Some(reason)) = self.banned_pubkeys.get(pubkey) {
let ban_info = BanInfo {
pubkey: pubkey.to_owned(),
reason: String::from_utf8(reason.to_vec()).unwrap(),
};
return Ok(Some(ban_info));
}
return Ok(None);
}
Ok(None)
}
}
#[cfg(test)]
mod tests {}