use super::{OwnerSource, Profile, ProfileChanges, ProfileImageChanges, StoreError, Undo}; use chrono::{DateTime, Utc}; use sled::{Db, Transactional, Tree}; use std::io::Cursor; use uuid::Uuid; #[derive(Clone, Debug)] pub struct Store { profile_tree: Tree, handle_tree: Tree, owner_tree: Tree, owner_created_tree: Tree, created_tree: Tree, count_tree: Tree, } #[derive(Clone, Debug, serde::Deserialize, serde::Serialize)] enum StoredOwnerSource { Local(Uuid), Remote(String), } #[derive(Debug, serde::Deserialize, serde::Serialize)] struct StoredProfile<'a> { id: Uuid, owner_source: StoredOwnerSource, handle: &'a str, domain: &'a str, #[serde(skip_serializing_if = "Option::is_none")] display_name: Option, #[serde(skip_serializing_if = "Option::is_none")] description: Option, #[serde(skip_serializing_if = "Option::is_none")] icon: Option, #[serde(skip_serializing_if = "Option::is_none")] banner: Option, published: DateTime, login_required: bool, created_at: DateTime, updated_at: DateTime, #[serde(skip_serializing_if = "Option::is_none")] suspended_at: Option>, } impl Store { pub(super) fn build(db: &Db) -> Result { Ok(Store { profile_tree: db.open_tree("profiles/profiles")?, handle_tree: db.open_tree("profiles/profiles/handles")?, owner_tree: db.open_tree("profiles/profiles/owner")?, owner_created_tree: db.open_tree("/profiles/profiles/owner/created")?, created_tree: db.open_tree("/profiles/profiles/created")?, count_tree: db.open_tree("/profiles/profiles/count")?, }) } pub(crate) fn create( &self, source: OwnerSource, handle: &str, domain: &str, published: DateTime, ) -> Result { let mut id; let mut stored_profile; let now = Utc::now().into(); let mut stored_profile_vec = vec![]; let stored_source = match &source { OwnerSource::Local(id) => StoredOwnerSource::Local(*id), OwnerSource::Remote(s) => StoredOwnerSource::Remote(s.clone()), }; while { stored_profile_vec.clear(); let writer = Cursor::new(&mut stored_profile_vec); id = Uuid::new_v4(); stored_profile = StoredProfile { id, owner_source: stored_source.clone(), handle, domain, display_name: None, description: None, icon: None, banner: None, published, login_required: true, created_at: now, updated_at: now, suspended_at: None, }; serde_json::to_writer(writer, &stored_profile)?; self.profile_tree .compare_and_swap( id_profile_key(id), None as Option<&[u8]>, Some(stored_profile_vec.as_slice()), )? .is_err() } {} let source = source.clone(); // ensure handle uniqueness match self.handle_tree.compare_and_swap( handle_id_key(handle, domain), None as Option<&[u8]>, Some(id.to_string().as_bytes()), ) { Ok(Ok(_)) => (), Ok(Err(_)) => { self.profile_tree.remove(id_profile_key(id))?; return Err(StoreError::DoubleStore); } Err(e) => { self.profile_tree.remove(id_profile_key(id))?; return Err(e.into()); } } let res = [ &self.owner_tree, &self.created_tree, &self.owner_created_tree, &self.count_tree, ] .transaction(move |trees| { let owner_tree = &trees[0]; let created_tree = &trees[1]; let owner_created_tree = &trees[2]; let count_tree = &trees[3]; owner_tree.insert(owner_key(&source, id).as_bytes(), id.to_string().as_bytes())?; created_tree.insert( created_profile_key(now, id).as_bytes(), id.to_string().as_bytes(), )?; owner_created_tree.insert( owner_created_profile_key(&source, now, id).as_bytes(), id.to_string().as_bytes(), )?; super::count(count_tree, &owner_profile_count_key(&source), |c| { c.saturating_add(1) })?; Ok(()) }); if let Err(e) = res { [&self.profile_tree, &self.handle_tree].transaction(move |trees| { let profile_tree = &trees[0]; let handle_tree = &trees[1]; profile_tree.remove(id_profile_key(id).as_bytes())?; handle_tree.remove(handle_id_key(handle, domain).as_bytes())?; Ok(()) })?; return Err(e.into()); } Ok(stored_profile.into()) } pub fn count(&self, source: OwnerSource) -> Result { match self.count_tree.get(owner_profile_count_key(&source))? { Some(ivec) => Ok(String::from_utf8_lossy(&ivec) .parse::() .expect("Count is valid")), None => Ok(0), } } pub(crate) fn update(&self, profile_changes: &ProfileChanges) -> Result { let stored_profile_ivec = match self .profile_tree .get(id_profile_key(profile_changes.id).as_bytes())? { Some(ivec) => ivec, None => return Err(StoreError::Missing), }; let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?; if let Some(login_required) = profile_changes.login_required { stored_profile.login_required = login_required; } if let Some(display_name) = profile_changes.display_name.as_ref() { stored_profile.display_name = Some(display_name.clone()); } if let Some(description) = profile_changes.description.as_ref() { stored_profile.description = Some(description.clone()); } stored_profile.updated_at = Utc::now().into(); let stored_profile_vec = serde_json::to_vec(&stored_profile)?; if self .profile_tree .compare_and_swap( id_profile_key(profile_changes.id).as_bytes(), Some(&stored_profile_ivec), Some(stored_profile_vec), )? .is_err() { return Err(StoreError::DoubleStore); } Ok(stored_profile.into()) } pub(crate) fn update_images( &self, profile_image_changes: &ProfileImageChanges, ) -> Result { let stored_profile_ivec = match self .profile_tree .get(id_profile_key(profile_image_changes.id).as_bytes())? { Some(ivec) => ivec, None => return Err(StoreError::Missing), }; let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?; if let Some(icon) = profile_image_changes.icon { stored_profile.icon = Some(icon); } if let Some(banner) = profile_image_changes.banner { stored_profile.banner = Some(banner); } stored_profile.updated_at = Utc::now().into(); let stored_profile_vec = serde_json::to_vec(&stored_profile)?; if self .profile_tree .compare_and_swap( id_profile_key(profile_image_changes.id).as_bytes(), Some(&stored_profile_ivec), Some(stored_profile_vec), )? .is_err() { return Err(StoreError::DoubleStore); } Ok(stored_profile.into()) } pub fn for_local(&self, id: Uuid) -> impl DoubleEndedIterator { self.for_owner(OwnerSource::Local(id)) } fn for_owner(&self, owner_source: OwnerSource) -> impl DoubleEndedIterator { self.owner_tree .scan_prefix(owner_prefix(&owner_source)) .values() .filter_map(|res| res.ok()) .filter_map(|ivec| { let id_str = String::from_utf8_lossy(&ivec); id_str.parse::().ok() }) } fn date_range( &self, range: impl std::ops::RangeBounds, ) -> impl DoubleEndedIterator where K: AsRef<[u8]>, { self.owner_created_tree .range(range) .values() .filter_map(|res| res.ok()) .filter_map(move |ivec| { let id_str = String::from_utf8_lossy(&ivec); id_str.parse::().ok() }) } pub fn newer_than(&self, id: Uuid) -> impl DoubleEndedIterator { let this = self.clone(); self.profile_tree .get(id_profile_key(id)) .ok() .and_then(|opt| opt) .and_then(|stored_profile_ivec| { let stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec).ok()?; Some(( stored_profile.owner_source.into(), stored_profile.created_at, )) }) .into_iter() .flat_map(move |(source, created_at)| { let range_start = owner_created_profile_range_start(&source, created_at); let range_start = range_start.as_bytes().to_vec(); this.date_range(range_start..) }) } pub fn older_than(&self, id: Uuid) -> impl DoubleEndedIterator { let this = self.clone(); self.profile_tree .get(id_profile_key(id)) .ok() .and_then(|opt| opt) .and_then(|stored_profile_ivec| { let stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec).ok()?; Some(( stored_profile.owner_source.into(), stored_profile.created_at, )) }) .into_iter() .flat_map(move |(source, created_at)| { let range_end = owner_created_profile_range_start(&source, created_at); let range_end = range_end.as_bytes().to_vec(); this.date_range(..range_end) }) } pub fn is_local(&self, id: Uuid) -> Result, StoreError> { if let Some(profile) = self.by_id(id)? { return Ok(Some(profile.owner_source.is_local())); } Ok(None) } pub fn by_id(&self, id: Uuid) -> Result, StoreError> { let stored_profile_ivec = match self.profile_tree.get(id_profile_key(id).as_bytes())? { Some(ivec) => ivec, None => return Ok(None), }; let stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?; Ok(Some(stored_profile.into())) } pub fn by_handle(&self, handle: &str, domain: &str) -> Result, StoreError> { let id_ivec = match self.handle_tree.get(handle_id_key(handle, domain))? { Some(ivec) => ivec, None => return Ok(None), }; let id_str = String::from_utf8_lossy(&id_ivec); Ok(id_str.parse().ok()) } pub(crate) fn suspend(&self, profile_id: Uuid) -> Result>, StoreError> { let stored_profile_ivec = match self .profile_tree .get(id_profile_key(profile_id).as_bytes())? { Some(ivec) => ivec, None => return Ok(None), }; let now = Utc::now(); let mut stored_profile: StoredProfile = serde_json::from_slice(&stored_profile_ivec)?; stored_profile.banner = None; stored_profile.icon = None; stored_profile.description = None; stored_profile.display_name = None; stored_profile.suspended_at = Some(now); stored_profile.updated_at = now; let stored_profile_vec = serde_json::to_vec(&stored_profile)?; if self .profile_tree .compare_and_swap( id_profile_key(profile_id).as_bytes(), Some(&stored_profile_ivec), Some(stored_profile_vec), )? .is_err() { return Err(StoreError::DoubleStore); } Ok(Some(Undo(stored_profile.into()))) } } // Used to map id -> Profile fn id_profile_key(id: Uuid) -> String { format!("/profile/{}/data", id) } // Used to map handle -> id fn handle_id_key(handle: &str, domain: &str) -> String { format!("/handle/{}:{}/profile", handle, domain) } // Used to map owner_id -> id fn owner_key(source: &OwnerSource, id: Uuid) -> String { format!("/owner/{}/profile/{}", source, id) } fn owner_prefix(source: &OwnerSource) -> String { format!("/owner/{}/profile", source) } // Used to fetch profiles for a given owner in a user-recognizalbe order fn owner_created_profile_key(owner: &OwnerSource, created_at: DateTime, id: Uuid) -> String { format!( "/owner/{}/created/{}/profile/{}", owner, created_at.to_rfc3339(), id ) } fn owner_created_profile_range_start(owner: &OwnerSource, created_at: DateTime) -> String { format!("/owner/{}/created/{}", owner, created_at.to_rfc3339(),) } // Used to map created_at -> id fn created_profile_key(created_at: DateTime, id: Uuid) -> String { format!("/created/{}/profile/{}", created_at, id) } fn owner_profile_count_key(owner: &OwnerSource) -> String { format!("/owner/{}/count", owner) } impl From for OwnerSource { fn from(os: StoredOwnerSource) -> Self { match os { StoredOwnerSource::Local(id) => OwnerSource::Local(id), StoredOwnerSource::Remote(s) => OwnerSource::Remote(s), } } } impl<'a> From> for Profile { fn from(sp: StoredProfile<'a>) -> Self { Profile { id: sp.id, owner_source: sp.owner_source.into(), handle: sp.handle.to_owned(), domain: sp.domain.to_owned(), display_name: sp.display_name, description: sp.description, icon: sp.icon, banner: sp.banner, published: sp.published, login_required: sp.login_required, suspended: sp.suspended_at.is_some(), } } }