dissolve key_value/*
Signed-off-by: Jason Volk <jason@zemos.net>
This commit is contained in:
parent
3122648767
commit
b94045a468
88 changed files with 4556 additions and 4751 deletions
|
@ -1,6 +1,6 @@
|
|||
use ruma::{OwnedRoomAliasId, OwnedRoomId, RoomAliasId, RoomId};
|
||||
use ruma::{api::client::error::ErrorKind, OwnedRoomAliasId, OwnedRoomId, RoomAliasId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
/// Creates or updates the alias to the given room id.
|
||||
|
@ -20,3 +20,75 @@ pub trait Data: Send + Sync {
|
|||
/// Returns all local aliases on the server
|
||||
fn all_local_aliases<'a>(&'a self) -> Box<dyn Iterator<Item = Result<(OwnedRoomId, String)>> + 'a>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn set_alias(&self, alias: &RoomAliasId, room_id: &RoomId) -> Result<()> {
|
||||
self.alias_roomid
|
||||
.insert(alias.alias().as_bytes(), room_id.as_bytes())?;
|
||||
let mut aliasid = room_id.as_bytes().to_vec();
|
||||
aliasid.push(0xFF);
|
||||
aliasid.extend_from_slice(&services().globals.next_count()?.to_be_bytes());
|
||||
self.aliasid_alias.insert(&aliasid, alias.as_bytes())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn remove_alias(&self, alias: &RoomAliasId) -> Result<()> {
|
||||
if let Some(room_id) = self.alias_roomid.get(alias.alias().as_bytes())? {
|
||||
let mut prefix = room_id;
|
||||
prefix.push(0xFF);
|
||||
|
||||
for (key, _) in self.aliasid_alias.scan_prefix(prefix) {
|
||||
self.aliasid_alias.remove(&key)?;
|
||||
}
|
||||
self.alias_roomid.remove(alias.alias().as_bytes())?;
|
||||
} else {
|
||||
return Err(Error::BadRequest(ErrorKind::NotFound, "Alias does not exist."));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn resolve_local_alias(&self, alias: &RoomAliasId) -> Result<Option<OwnedRoomId>> {
|
||||
self.alias_roomid
|
||||
.get(alias.alias().as_bytes())?
|
||||
.map(|bytes| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Room ID in alias_roomid is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in alias_roomid is invalid."))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn local_aliases_for_room<'a>(
|
||||
&'a self, room_id: &RoomId,
|
||||
) -> Box<dyn Iterator<Item = Result<OwnedRoomAliasId>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(self.aliasid_alias.scan_prefix(prefix).map(|(_, bytes)| {
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid alias bytes in aliasid_alias."))?
|
||||
.try_into()
|
||||
.map_err(|_| Error::bad_database("Invalid alias in aliasid_alias."))
|
||||
}))
|
||||
}
|
||||
|
||||
fn all_local_aliases<'a>(&'a self) -> Box<dyn Iterator<Item = Result<(OwnedRoomId, String)>> + 'a> {
|
||||
Box::new(
|
||||
self.alias_roomid
|
||||
.iter()
|
||||
.map(|(room_alias_bytes, room_id_bytes)| {
|
||||
let room_alias_localpart = utils::string_from_bytes(&room_alias_bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid alias bytes in aliasid_alias."))?;
|
||||
|
||||
let room_id = utils::string_from_bytes(&room_id_bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid room_id bytes in aliasid_alias."))?
|
||||
.try_into()
|
||||
.map_err(|_| Error::bad_database("Invalid room_id in aliasid_alias."))?;
|
||||
|
||||
Ok((room_id, room_alias_localpart))
|
||||
}),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{OwnedRoomAliasId, OwnedRoomId, RoomAliasId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
|
|
|
@ -1,8 +1,64 @@
|
|||
use std::sync::Arc;
|
||||
use std::{mem::size_of, sync::Arc};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{utils, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn get_cached_eventid_authchain(&self, shorteventid: &[u64]) -> Result<Option<Arc<[u64]>>>;
|
||||
fn cache_auth_chain(&self, shorteventid: Vec<u64>, auth_chain: Arc<[u64]>) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn get_cached_eventid_authchain(&self, key: &[u64]) -> Result<Option<Arc<[u64]>>> {
|
||||
// Check RAM cache
|
||||
if let Some(result) = self.auth_chain_cache.lock().unwrap().get_mut(key) {
|
||||
return Ok(Some(Arc::clone(result)));
|
||||
}
|
||||
|
||||
// We only save auth chains for single events in the db
|
||||
if key.len() == 1 {
|
||||
// Check DB cache
|
||||
let chain = self
|
||||
.shorteventid_authchain
|
||||
.get(&key[0].to_be_bytes())?
|
||||
.map(|chain| {
|
||||
chain
|
||||
.chunks_exact(size_of::<u64>())
|
||||
.map(|chunk| utils::u64_from_bytes(chunk).expect("byte length is correct"))
|
||||
.collect::<Arc<[u64]>>()
|
||||
});
|
||||
|
||||
if let Some(chain) = chain {
|
||||
// Cache in RAM
|
||||
self.auth_chain_cache
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(vec![key[0]], Arc::clone(&chain));
|
||||
|
||||
return Ok(Some(chain));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn cache_auth_chain(&self, key: Vec<u64>, auth_chain: Arc<[u64]>) -> Result<()> {
|
||||
// Only persist single events in db
|
||||
if key.len() == 1 {
|
||||
self.shorteventid_authchain.insert(
|
||||
&key[0].to_be_bytes(),
|
||||
&auth_chain
|
||||
.iter()
|
||||
.flat_map(|s| s.to_be_bytes().to_vec())
|
||||
.collect::<Vec<u8>>(),
|
||||
)?;
|
||||
}
|
||||
|
||||
// Cache in RAM
|
||||
self.auth_chain_cache
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(key, auth_chain);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -4,7 +4,7 @@ use std::{
|
|||
sync::Arc,
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{api::client::error::ErrorKind, EventId, RoomId};
|
||||
use tracing::{debug, error, trace, warn};
|
||||
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use ruma::{OwnedRoomId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
/// Adds the room to the public room directory
|
||||
|
@ -15,3 +15,23 @@ pub trait Data: Send + Sync {
|
|||
/// Returns the unsorted public room directory
|
||||
fn public_rooms<'a>(&'a self) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn set_public(&self, room_id: &RoomId) -> Result<()> { self.publicroomids.insert(room_id.as_bytes(), &[]) }
|
||||
|
||||
fn set_not_public(&self, room_id: &RoomId) -> Result<()> { self.publicroomids.remove(room_id.as_bytes()) }
|
||||
|
||||
fn is_public_room(&self, room_id: &RoomId) -> Result<bool> {
|
||||
Ok(self.publicroomids.get(room_id.as_bytes())?.is_some())
|
||||
}
|
||||
|
||||
fn public_rooms<'a>(&'a self) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a> {
|
||||
Box::new(self.publicroomids.iter().map(|(bytes, _)| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Room ID in publicroomids is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in publicroomids is invalid."))
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{OwnedRoomId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use ruma::{DeviceId, RoomId, UserId};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn lazy_load_was_sent_before(
|
||||
|
@ -14,3 +14,53 @@ pub trait Data: Send + Sync {
|
|||
|
||||
fn lazy_load_reset(&self, user_id: &UserId, device_id: &DeviceId, room_id: &RoomId) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn lazy_load_was_sent_before(
|
||||
&self, user_id: &UserId, device_id: &DeviceId, room_id: &RoomId, ll_user: &UserId,
|
||||
) -> Result<bool> {
|
||||
let mut key = user_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(device_id.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(room_id.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(ll_user.as_bytes());
|
||||
Ok(self.lazyloadedids.get(&key)?.is_some())
|
||||
}
|
||||
|
||||
fn lazy_load_confirm_delivery(
|
||||
&self, user_id: &UserId, device_id: &DeviceId, room_id: &RoomId,
|
||||
confirmed_user_ids: &mut dyn Iterator<Item = &UserId>,
|
||||
) -> Result<()> {
|
||||
let mut prefix = user_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
prefix.extend_from_slice(device_id.as_bytes());
|
||||
prefix.push(0xFF);
|
||||
prefix.extend_from_slice(room_id.as_bytes());
|
||||
prefix.push(0xFF);
|
||||
|
||||
for ll_id in confirmed_user_ids {
|
||||
let mut key = prefix.clone();
|
||||
key.extend_from_slice(ll_id.as_bytes());
|
||||
self.lazyloadedids.insert(&key, &[])?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn lazy_load_reset(&self, user_id: &UserId, device_id: &DeviceId, room_id: &RoomId) -> Result<()> {
|
||||
let mut prefix = user_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
prefix.extend_from_slice(device_id.as_bytes());
|
||||
prefix.push(0xFF);
|
||||
prefix.extend_from_slice(room_id.as_bytes());
|
||||
prefix.push(0xFF);
|
||||
|
||||
for (key, _) in self.lazyloadedids.scan_prefix(prefix) {
|
||||
self.lazyloadedids.remove(&key)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -4,7 +4,7 @@ use std::{
|
|||
sync::Arc,
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{DeviceId, OwnedDeviceId, OwnedRoomId, OwnedUserId, RoomId, UserId};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
|
|
|
@ -1,6 +1,7 @@
|
|||
use ruma::{OwnedRoomId, RoomId};
|
||||
use tracing::error;
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn exists(&self, room_id: &RoomId) -> Result<bool>;
|
||||
|
@ -11,3 +12,75 @@ pub trait Data: Send + Sync {
|
|||
fn ban_room(&self, room_id: &RoomId, banned: bool) -> Result<()>;
|
||||
fn list_banned_rooms<'a>(&'a self) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn exists(&self, room_id: &RoomId) -> Result<bool> {
|
||||
let prefix = match services().rooms.short.get_shortroomid(room_id)? {
|
||||
Some(b) => b.to_be_bytes().to_vec(),
|
||||
None => return Ok(false),
|
||||
};
|
||||
|
||||
// Look for PDUs in that room.
|
||||
Ok(self
|
||||
.pduid_pdu
|
||||
.iter_from(&prefix, false)
|
||||
.next()
|
||||
.filter(|(k, _)| k.starts_with(&prefix))
|
||||
.is_some())
|
||||
}
|
||||
|
||||
fn iter_ids<'a>(&'a self) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a> {
|
||||
Box::new(self.roomid_shortroomid.iter().map(|(bytes, _)| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Room ID in publicroomids is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in roomid_shortroomid is invalid."))
|
||||
}))
|
||||
}
|
||||
|
||||
fn is_disabled(&self, room_id: &RoomId) -> Result<bool> {
|
||||
Ok(self.disabledroomids.get(room_id.as_bytes())?.is_some())
|
||||
}
|
||||
|
||||
fn disable_room(&self, room_id: &RoomId, disabled: bool) -> Result<()> {
|
||||
if disabled {
|
||||
self.disabledroomids.insert(room_id.as_bytes(), &[])?;
|
||||
} else {
|
||||
self.disabledroomids.remove(room_id.as_bytes())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_banned(&self, room_id: &RoomId) -> Result<bool> { Ok(self.bannedroomids.get(room_id.as_bytes())?.is_some()) }
|
||||
|
||||
fn ban_room(&self, room_id: &RoomId, banned: bool) -> Result<()> {
|
||||
if banned {
|
||||
self.bannedroomids.insert(room_id.as_bytes(), &[])?;
|
||||
} else {
|
||||
self.bannedroomids.remove(room_id.as_bytes())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn list_banned_rooms<'a>(&'a self) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a> {
|
||||
Box::new(self.bannedroomids.iter().map(
|
||||
|(room_id_bytes, _ /* non-banned rooms should not be in this table */)| {
|
||||
let room_id = utils::string_from_bytes(&room_id_bytes)
|
||||
.map_err(|e| {
|
||||
error!("Invalid room_id bytes in bannedroomids: {e}");
|
||||
Error::bad_database("Invalid room_id in bannedroomids.")
|
||||
})?
|
||||
.try_into()
|
||||
.map_err(|e| {
|
||||
error!("Invalid room_id in bannedroomids: {e}");
|
||||
Error::bad_database("Invalid room_id in bannedroomids")
|
||||
})?;
|
||||
|
||||
Ok(room_id)
|
||||
},
|
||||
))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{OwnedRoomId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
|
|
|
@ -19,27 +19,6 @@ pub mod timeline;
|
|||
pub mod typing;
|
||||
pub mod user;
|
||||
|
||||
pub trait Data:
|
||||
alias::Data
|
||||
+ auth_chain::Data
|
||||
+ directory::Data
|
||||
+ lazy_loading::Data
|
||||
+ metadata::Data
|
||||
+ outlier::Data
|
||||
+ pdu_metadata::Data
|
||||
+ read_receipt::Data
|
||||
+ search::Data
|
||||
+ short::Data
|
||||
+ state::Data
|
||||
+ state_accessor::Data
|
||||
+ state_cache::Data
|
||||
+ state_compressor::Data
|
||||
+ timeline::Data
|
||||
+ threads::Data
|
||||
+ user::Data
|
||||
{
|
||||
}
|
||||
|
||||
pub struct Service {
|
||||
pub alias: alias::Service,
|
||||
pub auth_chain: auth_chain::Service,
|
||||
|
|
|
@ -1,9 +1,34 @@
|
|||
use ruma::{CanonicalJsonObject, EventId};
|
||||
|
||||
use crate::{PduEvent, Result};
|
||||
use crate::{Error, KeyValueDatabase, PduEvent, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn get_outlier_pdu_json(&self, event_id: &EventId) -> Result<Option<CanonicalJsonObject>>;
|
||||
fn get_outlier_pdu(&self, event_id: &EventId) -> Result<Option<PduEvent>>;
|
||||
fn add_pdu_outlier(&self, event_id: &EventId, pdu: &CanonicalJsonObject) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn get_outlier_pdu_json(&self, event_id: &EventId) -> Result<Option<CanonicalJsonObject>> {
|
||||
self.eventid_outlierpdu
|
||||
.get(event_id.as_bytes())?
|
||||
.map_or(Ok(None), |pdu| {
|
||||
serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db."))
|
||||
})
|
||||
}
|
||||
|
||||
fn get_outlier_pdu(&self, event_id: &EventId) -> Result<Option<PduEvent>> {
|
||||
self.eventid_outlierpdu
|
||||
.get(event_id.as_bytes())?
|
||||
.map_or(Ok(None), |pdu| {
|
||||
serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db."))
|
||||
})
|
||||
}
|
||||
|
||||
fn add_pdu_outlier(&self, event_id: &EventId, pdu: &CanonicalJsonObject) -> Result<()> {
|
||||
self.eventid_outlierpdu.insert(
|
||||
event_id.as_bytes(),
|
||||
&serde_json::to_vec(&pdu).expect("CanonicalJsonObject is valid"),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{CanonicalJsonObject, EventId};
|
||||
|
||||
use crate::{PduEvent, Result};
|
||||
|
|
|
@ -1,8 +1,8 @@
|
|||
use std::sync::Arc;
|
||||
use std::{mem, sync::Arc};
|
||||
|
||||
use ruma::{EventId, RoomId, UserId};
|
||||
|
||||
use crate::{PduCount, PduEvent, Result};
|
||||
use crate::{services, utils, Error, KeyValueDatabase, PduCount, PduEvent, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn add_relation(&self, from: u64, to: u64) -> Result<()>;
|
||||
|
@ -15,3 +15,77 @@ pub trait Data: Send + Sync {
|
|||
fn mark_event_soft_failed(&self, event_id: &EventId) -> Result<()>;
|
||||
fn is_event_soft_failed(&self, event_id: &EventId) -> Result<bool>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn add_relation(&self, from: u64, to: u64) -> Result<()> {
|
||||
let mut key = to.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(&from.to_be_bytes());
|
||||
self.tofrom_relation.insert(&key, &[])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn relations_until<'a>(
|
||||
&'a self, user_id: &'a UserId, shortroomid: u64, target: u64, until: PduCount,
|
||||
) -> Result<Box<dyn Iterator<Item = Result<(PduCount, PduEvent)>> + 'a>> {
|
||||
let prefix = target.to_be_bytes().to_vec();
|
||||
let mut current = prefix.clone();
|
||||
|
||||
let count_raw = match until {
|
||||
PduCount::Normal(x) => x.saturating_sub(1),
|
||||
PduCount::Backfilled(x) => {
|
||||
current.extend_from_slice(&0_u64.to_be_bytes());
|
||||
u64::MAX.saturating_sub(x).saturating_sub(1)
|
||||
},
|
||||
};
|
||||
current.extend_from_slice(&count_raw.to_be_bytes());
|
||||
|
||||
Ok(Box::new(
|
||||
self.tofrom_relation
|
||||
.iter_from(¤t, true)
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix))
|
||||
.map(move |(tofrom, _data)| {
|
||||
let from = utils::u64_from_bytes(&tofrom[(mem::size_of::<u64>())..])
|
||||
.map_err(|_| Error::bad_database("Invalid count in tofrom_relation."))?;
|
||||
|
||||
let mut pduid = shortroomid.to_be_bytes().to_vec();
|
||||
pduid.extend_from_slice(&from.to_be_bytes());
|
||||
|
||||
let mut pdu = services()
|
||||
.rooms
|
||||
.timeline
|
||||
.get_pdu_from_id(&pduid)?
|
||||
.ok_or_else(|| Error::bad_database("Pdu in tofrom_relation is invalid."))?;
|
||||
if pdu.sender != user_id {
|
||||
pdu.remove_transaction_id()?;
|
||||
}
|
||||
Ok((PduCount::Normal(from), pdu))
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn mark_as_referenced(&self, room_id: &RoomId, event_ids: &[Arc<EventId>]) -> Result<()> {
|
||||
for prev in event_ids {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.extend_from_slice(prev.as_bytes());
|
||||
self.referencedevents.insert(&key, &[])?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_event_referenced(&self, room_id: &RoomId, event_id: &EventId) -> Result<bool> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.extend_from_slice(event_id.as_bytes());
|
||||
Ok(self.referencedevents.get(&key)?.is_some())
|
||||
}
|
||||
|
||||
fn mark_event_soft_failed(&self, event_id: &EventId) -> Result<()> {
|
||||
self.softfailedeventids.insert(event_id.as_bytes(), &[])
|
||||
}
|
||||
|
||||
fn is_event_soft_failed(&self, event_id: &EventId) -> Result<bool> {
|
||||
self.softfailedeventids
|
||||
.get(event_id.as_bytes())
|
||||
.map(|o| o.is_some())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{
|
||||
api::{client::relations::get_relating_events, Direction},
|
||||
events::{relation::RelationType, TimelineEventType},
|
||||
|
|
|
@ -1,10 +1,12 @@
|
|||
use std::mem;
|
||||
|
||||
use ruma::{
|
||||
events::{receipt::ReceiptEvent, AnySyncEphemeralRoomEvent},
|
||||
serde::Raw,
|
||||
OwnedUserId, RoomId, UserId,
|
||||
CanonicalJsonObject, OwnedUserId, RoomId, UserId,
|
||||
};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
type AnySyncEphemeralRoomEventIter<'a> =
|
||||
Box<dyn Iterator<Item = Result<(OwnedUserId, u64, Raw<AnySyncEphemeralRoomEvent>)>> + 'a>;
|
||||
|
@ -32,3 +34,118 @@ pub trait Data: Send + Sync {
|
|||
#[allow(dead_code)]
|
||||
fn last_privateread_update(&self, user_id: &UserId, room_id: &RoomId) -> Result<u64>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn readreceipt_update(&self, user_id: &UserId, room_id: &RoomId, event: ReceiptEvent) -> Result<()> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
let mut last_possible_key = prefix.clone();
|
||||
last_possible_key.extend_from_slice(&u64::MAX.to_be_bytes());
|
||||
|
||||
// Remove old entry
|
||||
if let Some((old, _)) = self
|
||||
.readreceiptid_readreceipt
|
||||
.iter_from(&last_possible_key, true)
|
||||
.take_while(|(key, _)| key.starts_with(&prefix))
|
||||
.find(|(key, _)| {
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element")
|
||||
== user_id.as_bytes()
|
||||
}) {
|
||||
// This is the old room_latest
|
||||
self.readreceiptid_readreceipt.remove(&old)?;
|
||||
}
|
||||
|
||||
let mut room_latest_id = prefix;
|
||||
room_latest_id.extend_from_slice(&services().globals.next_count()?.to_be_bytes());
|
||||
room_latest_id.push(0xFF);
|
||||
room_latest_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.readreceiptid_readreceipt.insert(
|
||||
&room_latest_id,
|
||||
&serde_json::to_vec(&event).expect("EduEvent::to_string always works"),
|
||||
)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn readreceipts_since<'a>(
|
||||
&'a self, room_id: &RoomId, since: u64,
|
||||
) -> Box<dyn Iterator<Item = Result<(OwnedUserId, u64, Raw<AnySyncEphemeralRoomEvent>)>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
let prefix2 = prefix.clone();
|
||||
|
||||
let mut first_possible_edu = prefix.clone();
|
||||
first_possible_edu.extend_from_slice(&(since.saturating_add(1)).to_be_bytes()); // +1 so we don't send the event at since
|
||||
|
||||
Box::new(
|
||||
self.readreceiptid_readreceipt
|
||||
.iter_from(&first_possible_edu, false)
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix2))
|
||||
.map(move |(k, v)| {
|
||||
let count = utils::u64_from_bytes(&k[prefix.len()..prefix.len() + mem::size_of::<u64>()])
|
||||
.map_err(|_| Error::bad_database("Invalid readreceiptid count in db."))?;
|
||||
let user_id = UserId::parse(
|
||||
utils::string_from_bytes(&k[prefix.len() + mem::size_of::<u64>() + 1..])
|
||||
.map_err(|_| Error::bad_database("Invalid readreceiptid userid bytes in db."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Invalid readreceiptid userid in db."))?;
|
||||
|
||||
let mut json = serde_json::from_slice::<CanonicalJsonObject>(&v)
|
||||
.map_err(|_| Error::bad_database("Read receipt in roomlatestid_roomlatest is invalid json."))?;
|
||||
json.remove("room_id");
|
||||
|
||||
Ok((
|
||||
user_id,
|
||||
count,
|
||||
Raw::from_json(serde_json::value::to_raw_value(&json).expect("json is valid raw value")),
|
||||
))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
fn private_read_set(&self, room_id: &RoomId, user_id: &UserId, count: u64) -> Result<()> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.roomuserid_privateread
|
||||
.insert(&key, &count.to_be_bytes())?;
|
||||
|
||||
self.roomuserid_lastprivatereadupdate
|
||||
.insert(&key, &services().globals.next_count()?.to_be_bytes())
|
||||
}
|
||||
|
||||
fn private_read_get(&self, room_id: &RoomId, user_id: &UserId) -> Result<Option<u64>> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.roomuserid_privateread
|
||||
.get(&key)?
|
||||
.map_or(Ok(None), |v| {
|
||||
Ok(Some(
|
||||
utils::u64_from_bytes(&v).map_err(|_| Error::bad_database("Invalid private read marker bytes"))?,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn last_privateread_update(&self, user_id: &UserId, room_id: &RoomId) -> Result<u64> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
Ok(self
|
||||
.roomuserid_lastprivatereadupdate
|
||||
.get(&key)?
|
||||
.map(|bytes| {
|
||||
utils::u64_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Count in roomuserid_lastprivatereadupdate is invalid."))
|
||||
})
|
||||
.transpose()?
|
||||
.unwrap_or(0))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{events::receipt::ReceiptEvent, serde::Raw, OwnedUserId, RoomId, UserId};
|
||||
|
||||
use crate::{services, Result};
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use ruma::RoomId;
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, KeyValueDatabase, Result};
|
||||
|
||||
type SearchPdusResult<'a> = Result<Option<(Box<dyn Iterator<Item = Vec<u8>> + 'a>, Vec<String>)>>;
|
||||
|
||||
|
@ -9,3 +9,62 @@ pub trait Data: Send + Sync {
|
|||
|
||||
fn search_pdus<'a>(&'a self, room_id: &RoomId, search_string: &str) -> SearchPdusResult<'a>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn index_pdu(&self, shortroomid: u64, pdu_id: &[u8], message_body: &str) -> Result<()> {
|
||||
let mut batch = message_body
|
||||
.split_terminator(|c: char| !c.is_alphanumeric())
|
||||
.filter(|s| !s.is_empty())
|
||||
.filter(|word| word.len() <= 50)
|
||||
.map(str::to_lowercase)
|
||||
.map(|word| {
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(word.as_bytes());
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(pdu_id); // TODO: currently we save the room id a second time here
|
||||
(key, Vec::new())
|
||||
});
|
||||
|
||||
self.tokenids.insert_batch(&mut batch)
|
||||
}
|
||||
|
||||
fn search_pdus<'a>(&'a self, room_id: &RoomId, search_string: &str) -> SearchPdusResult<'a> {
|
||||
let prefix = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortroomid(room_id)?
|
||||
.expect("room exists")
|
||||
.to_be_bytes()
|
||||
.to_vec();
|
||||
|
||||
let words: Vec<_> = search_string
|
||||
.split_terminator(|c: char| !c.is_alphanumeric())
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_lowercase)
|
||||
.collect();
|
||||
|
||||
let iterators = words.clone().into_iter().map(move |word| {
|
||||
let mut prefix2 = prefix.clone();
|
||||
prefix2.extend_from_slice(word.as_bytes());
|
||||
prefix2.push(0xFF);
|
||||
let prefix3 = prefix2.clone();
|
||||
|
||||
let mut last_possible_id = prefix2.clone();
|
||||
last_possible_id.extend_from_slice(&u64::MAX.to_be_bytes());
|
||||
|
||||
self.tokenids
|
||||
.iter_from(&last_possible_id, true) // Newest pdus first
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix2))
|
||||
.map(move |(key, _)| key[prefix3.len()..].to_vec())
|
||||
});
|
||||
|
||||
let Some(common_elements) = utils::common_elements(iterators, |a, b| {
|
||||
// We compare b with a because we reversed the iterator earlier
|
||||
b.cmp(a)
|
||||
}) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
Ok(Some((Box::new(common_elements), words)))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::RoomId;
|
||||
|
||||
use crate::Result;
|
||||
|
|
|
@ -1,8 +1,9 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use ruma::{events::StateEventType, EventId, RoomId};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn get_or_create_shorteventid(&self, event_id: &EventId) -> Result<u64>;
|
||||
|
@ -24,3 +25,161 @@ pub trait Data: Send + Sync {
|
|||
|
||||
fn get_or_create_shortroomid(&self, room_id: &RoomId) -> Result<u64>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn get_or_create_shorteventid(&self, event_id: &EventId) -> Result<u64> {
|
||||
let short = if let Some(shorteventid) = self.eventid_shorteventid.get(event_id.as_bytes())? {
|
||||
utils::u64_from_bytes(&shorteventid).map_err(|_| Error::bad_database("Invalid shorteventid in db."))?
|
||||
} else {
|
||||
let shorteventid = services().globals.next_count()?;
|
||||
self.eventid_shorteventid
|
||||
.insert(event_id.as_bytes(), &shorteventid.to_be_bytes())?;
|
||||
self.shorteventid_eventid
|
||||
.insert(&shorteventid.to_be_bytes(), event_id.as_bytes())?;
|
||||
shorteventid
|
||||
};
|
||||
|
||||
Ok(short)
|
||||
}
|
||||
|
||||
fn multi_get_or_create_shorteventid(&self, event_ids: &[&EventId]) -> Result<Vec<u64>> {
|
||||
let mut ret: Vec<u64> = Vec::with_capacity(event_ids.len());
|
||||
let keys = event_ids
|
||||
.iter()
|
||||
.map(|id| id.as_bytes())
|
||||
.collect::<Vec<&[u8]>>();
|
||||
for (i, short) in self
|
||||
.eventid_shorteventid
|
||||
.multi_get(&keys)?
|
||||
.iter()
|
||||
.enumerate()
|
||||
{
|
||||
match short {
|
||||
Some(short) => ret.push(
|
||||
utils::u64_from_bytes(short).map_err(|_| Error::bad_database("Invalid shorteventid in db."))?,
|
||||
),
|
||||
None => {
|
||||
let short = services().globals.next_count()?;
|
||||
self.eventid_shorteventid
|
||||
.insert(keys[i], &short.to_be_bytes())?;
|
||||
self.shorteventid_eventid
|
||||
.insert(&short.to_be_bytes(), keys[i])?;
|
||||
|
||||
debug_assert!(ret.len() == i, "position of result must match input");
|
||||
ret.push(short);
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ret)
|
||||
}
|
||||
|
||||
fn get_shortstatekey(&self, event_type: &StateEventType, state_key: &str) -> Result<Option<u64>> {
|
||||
let mut statekey_vec = event_type.to_string().as_bytes().to_vec();
|
||||
statekey_vec.push(0xFF);
|
||||
statekey_vec.extend_from_slice(state_key.as_bytes());
|
||||
|
||||
let short = self
|
||||
.statekey_shortstatekey
|
||||
.get(&statekey_vec)?
|
||||
.map(|shortstatekey| {
|
||||
utils::u64_from_bytes(&shortstatekey).map_err(|_| Error::bad_database("Invalid shortstatekey in db."))
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
Ok(short)
|
||||
}
|
||||
|
||||
fn get_or_create_shortstatekey(&self, event_type: &StateEventType, state_key: &str) -> Result<u64> {
|
||||
let mut statekey_vec = event_type.to_string().as_bytes().to_vec();
|
||||
statekey_vec.push(0xFF);
|
||||
statekey_vec.extend_from_slice(state_key.as_bytes());
|
||||
|
||||
let short = if let Some(shortstatekey) = self.statekey_shortstatekey.get(&statekey_vec)? {
|
||||
utils::u64_from_bytes(&shortstatekey).map_err(|_| Error::bad_database("Invalid shortstatekey in db."))?
|
||||
} else {
|
||||
let shortstatekey = services().globals.next_count()?;
|
||||
self.statekey_shortstatekey
|
||||
.insert(&statekey_vec, &shortstatekey.to_be_bytes())?;
|
||||
self.shortstatekey_statekey
|
||||
.insert(&shortstatekey.to_be_bytes(), &statekey_vec)?;
|
||||
shortstatekey
|
||||
};
|
||||
|
||||
Ok(short)
|
||||
}
|
||||
|
||||
fn get_eventid_from_short(&self, shorteventid: u64) -> Result<Arc<EventId>> {
|
||||
let bytes = self
|
||||
.shorteventid_eventid
|
||||
.get(&shorteventid.to_be_bytes())?
|
||||
.ok_or_else(|| Error::bad_database("Shorteventid does not exist"))?;
|
||||
|
||||
let event_id = EventId::parse_arc(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("EventID in shorteventid_eventid is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("EventId in shorteventid_eventid is invalid."))?;
|
||||
|
||||
Ok(event_id)
|
||||
}
|
||||
|
||||
fn get_statekey_from_short(&self, shortstatekey: u64) -> Result<(StateEventType, String)> {
|
||||
let bytes = self
|
||||
.shortstatekey_statekey
|
||||
.get(&shortstatekey.to_be_bytes())?
|
||||
.ok_or_else(|| Error::bad_database("Shortstatekey does not exist"))?;
|
||||
|
||||
let mut parts = bytes.splitn(2, |&b| b == 0xFF);
|
||||
let eventtype_bytes = parts.next().expect("split always returns one entry");
|
||||
let statekey_bytes = parts
|
||||
.next()
|
||||
.ok_or_else(|| Error::bad_database("Invalid statekey in shortstatekey_statekey."))?;
|
||||
|
||||
let event_type = StateEventType::from(utils::string_from_bytes(eventtype_bytes).map_err(|e| {
|
||||
warn!("Event type in shortstatekey_statekey is invalid: {}", e);
|
||||
Error::bad_database("Event type in shortstatekey_statekey is invalid.")
|
||||
})?);
|
||||
|
||||
let state_key = utils::string_from_bytes(statekey_bytes)
|
||||
.map_err(|_| Error::bad_database("Statekey in shortstatekey_statekey is invalid unicode."))?;
|
||||
|
||||
let result = (event_type, state_key);
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Returns (shortstatehash, already_existed)
|
||||
fn get_or_create_shortstatehash(&self, state_hash: &[u8]) -> Result<(u64, bool)> {
|
||||
Ok(if let Some(shortstatehash) = self.statehash_shortstatehash.get(state_hash)? {
|
||||
(
|
||||
utils::u64_from_bytes(&shortstatehash)
|
||||
.map_err(|_| Error::bad_database("Invalid shortstatehash in db."))?,
|
||||
true,
|
||||
)
|
||||
} else {
|
||||
let shortstatehash = services().globals.next_count()?;
|
||||
self.statehash_shortstatehash
|
||||
.insert(state_hash, &shortstatehash.to_be_bytes())?;
|
||||
(shortstatehash, false)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_shortroomid(&self, room_id: &RoomId) -> Result<Option<u64>> {
|
||||
self.roomid_shortroomid
|
||||
.get(room_id.as_bytes())?
|
||||
.map(|bytes| utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid shortroomid in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn get_or_create_shortroomid(&self, room_id: &RoomId) -> Result<u64> {
|
||||
Ok(if let Some(short) = self.roomid_shortroomid.get(room_id.as_bytes())? {
|
||||
utils::u64_from_bytes(&short).map_err(|_| Error::bad_database("Invalid shortroomid in db."))?
|
||||
} else {
|
||||
let short = services().globals.next_count()?;
|
||||
self.roomid_shortroomid
|
||||
.insert(room_id.as_bytes(), &short.to_be_bytes())?;
|
||||
short
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
mod data;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{events::StateEventType, EventId, RoomId};
|
||||
|
||||
use crate::Result;
|
||||
|
|
|
@ -3,7 +3,7 @@ use std::{collections::HashSet, sync::Arc};
|
|||
use ruma::{EventId, OwnedEventId, RoomId};
|
||||
use tokio::sync::MutexGuard;
|
||||
|
||||
use crate::Result;
|
||||
use crate::{utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
/// Returns the last state hash key added to the db for the given room.
|
||||
|
@ -31,3 +31,70 @@ pub trait Data: Send + Sync {
|
|||
_mutex_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room state mutex
|
||||
) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn get_room_shortstatehash(&self, room_id: &RoomId) -> Result<Option<u64>> {
|
||||
self.roomid_shortstatehash
|
||||
.get(room_id.as_bytes())?
|
||||
.map_or(Ok(None), |bytes| {
|
||||
Ok(Some(utils::u64_from_bytes(&bytes).map_err(|_| {
|
||||
Error::bad_database("Invalid shortstatehash in roomid_shortstatehash")
|
||||
})?))
|
||||
})
|
||||
}
|
||||
|
||||
fn set_room_state(
|
||||
&self,
|
||||
room_id: &RoomId,
|
||||
new_shortstatehash: u64,
|
||||
_mutex_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room state mutex
|
||||
) -> Result<()> {
|
||||
self.roomid_shortstatehash
|
||||
.insert(room_id.as_bytes(), &new_shortstatehash.to_be_bytes())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_event_state(&self, shorteventid: u64, shortstatehash: u64) -> Result<()> {
|
||||
self.shorteventid_shortstatehash
|
||||
.insert(&shorteventid.to_be_bytes(), &shortstatehash.to_be_bytes())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_forward_extremities(&self, room_id: &RoomId) -> Result<HashSet<Arc<EventId>>> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
self.roomid_pduleaves
|
||||
.scan_prefix(prefix)
|
||||
.map(|(_, bytes)| {
|
||||
EventId::parse_arc(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("EventID in roomid_pduleaves is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("EventId in roomid_pduleaves is invalid."))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn set_forward_extremities(
|
||||
&self,
|
||||
room_id: &RoomId,
|
||||
event_ids: Vec<OwnedEventId>,
|
||||
_mutex_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room state mutex
|
||||
) -> Result<()> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
for (key, _) in self.roomid_pduleaves.scan_prefix(prefix.clone()) {
|
||||
self.roomid_pduleaves.remove(&key)?;
|
||||
}
|
||||
|
||||
for event_id in event_ids {
|
||||
let mut key = prefix.clone();
|
||||
key.extend_from_slice(event_id.as_bytes());
|
||||
self.roomid_pduleaves.insert(&key, event_id.as_bytes())?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -4,7 +4,7 @@ use std::{
|
|||
sync::Arc,
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{
|
||||
api::client::error::ErrorKind,
|
||||
events::{
|
||||
|
|
|
@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc};
|
|||
use async_trait::async_trait;
|
||||
use ruma::{events::StateEventType, EventId, RoomId};
|
||||
|
||||
use crate::{PduEvent, Result};
|
||||
use crate::{services, utils, Error, KeyValueDatabase, PduEvent, Result};
|
||||
|
||||
#[async_trait]
|
||||
pub trait Data: Send + Sync {
|
||||
|
@ -46,3 +46,162 @@ pub trait Data: Send + Sync {
|
|||
&self, room_id: &RoomId, event_type: &StateEventType, state_key: &str,
|
||||
) -> Result<Option<Arc<PduEvent>>>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Data for KeyValueDatabase {
|
||||
#[allow(unused_qualifications)] // async traits
|
||||
async fn state_full_ids(&self, shortstatehash: u64) -> Result<HashMap<u64, Arc<EventId>>> {
|
||||
let full_state = services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.load_shortstatehash_info(shortstatehash)?
|
||||
.pop()
|
||||
.expect("there is always one layer")
|
||||
.1;
|
||||
let mut result = HashMap::new();
|
||||
let mut i: u8 = 0;
|
||||
for compressed in full_state.iter() {
|
||||
let parsed = services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.parse_compressed_state_event(compressed)?;
|
||||
result.insert(parsed.0, parsed.1);
|
||||
|
||||
i = i.wrapping_add(1);
|
||||
if i % 100 == 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
#[allow(unused_qualifications)] // async traits
|
||||
async fn state_full(&self, shortstatehash: u64) -> Result<HashMap<(StateEventType, String), Arc<PduEvent>>> {
|
||||
let full_state = services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.load_shortstatehash_info(shortstatehash)?
|
||||
.pop()
|
||||
.expect("there is always one layer")
|
||||
.1;
|
||||
|
||||
let mut result = HashMap::new();
|
||||
let mut i: u8 = 0;
|
||||
for compressed in full_state.iter() {
|
||||
let (_, eventid) = services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.parse_compressed_state_event(compressed)?;
|
||||
if let Some(pdu) = services().rooms.timeline.get_pdu(&eventid)? {
|
||||
result.insert(
|
||||
(
|
||||
pdu.kind.to_string().into(),
|
||||
pdu.state_key
|
||||
.as_ref()
|
||||
.ok_or_else(|| Error::bad_database("State event has no state key."))?
|
||||
.clone(),
|
||||
),
|
||||
pdu,
|
||||
);
|
||||
}
|
||||
|
||||
i = i.wrapping_add(1);
|
||||
if i % 100 == 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Returns a single PDU from `room_id` with key (`event_type`,
|
||||
/// `state_key`).
|
||||
fn state_get_id(
|
||||
&self, shortstatehash: u64, event_type: &StateEventType, state_key: &str,
|
||||
) -> Result<Option<Arc<EventId>>> {
|
||||
let Some(shortstatekey) = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortstatekey(event_type, state_key)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let full_state = services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.load_shortstatehash_info(shortstatehash)?
|
||||
.pop()
|
||||
.expect("there is always one layer")
|
||||
.1;
|
||||
Ok(full_state
|
||||
.iter()
|
||||
.find(|bytes| bytes.starts_with(&shortstatekey.to_be_bytes()))
|
||||
.and_then(|compressed| {
|
||||
services()
|
||||
.rooms
|
||||
.state_compressor
|
||||
.parse_compressed_state_event(compressed)
|
||||
.ok()
|
||||
.map(|(_, id)| id)
|
||||
}))
|
||||
}
|
||||
|
||||
/// Returns a single PDU from `room_id` with key (`event_type`,
|
||||
/// `state_key`).
|
||||
fn state_get(
|
||||
&self, shortstatehash: u64, event_type: &StateEventType, state_key: &str,
|
||||
) -> Result<Option<Arc<PduEvent>>> {
|
||||
self.state_get_id(shortstatehash, event_type, state_key)?
|
||||
.map_or(Ok(None), |event_id| services().rooms.timeline.get_pdu(&event_id))
|
||||
}
|
||||
|
||||
/// Returns the state hash for this pdu.
|
||||
fn pdu_shortstatehash(&self, event_id: &EventId) -> Result<Option<u64>> {
|
||||
self.eventid_shorteventid
|
||||
.get(event_id.as_bytes())?
|
||||
.map_or(Ok(None), |shorteventid| {
|
||||
self.shorteventid_shortstatehash
|
||||
.get(&shorteventid)?
|
||||
.map(|bytes| {
|
||||
utils::u64_from_bytes(&bytes).map_err(|_| {
|
||||
Error::bad_database("Invalid shortstatehash bytes in shorteventid_shortstatehash")
|
||||
})
|
||||
})
|
||||
.transpose()
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the full room state.
|
||||
#[allow(unused_qualifications)] // async traits
|
||||
async fn room_state_full(&self, room_id: &RoomId) -> Result<HashMap<(StateEventType, String), Arc<PduEvent>>> {
|
||||
if let Some(current_shortstatehash) = services().rooms.state.get_room_shortstatehash(room_id)? {
|
||||
self.state_full(current_shortstatehash).await
|
||||
} else {
|
||||
Ok(HashMap::new())
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a single PDU from `room_id` with key (`event_type`,
|
||||
/// `state_key`).
|
||||
fn room_state_get_id(
|
||||
&self, room_id: &RoomId, event_type: &StateEventType, state_key: &str,
|
||||
) -> Result<Option<Arc<EventId>>> {
|
||||
if let Some(current_shortstatehash) = services().rooms.state.get_room_shortstatehash(room_id)? {
|
||||
self.state_get_id(current_shortstatehash, event_type, state_key)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a single PDU from `room_id` with key (`event_type`,
|
||||
/// `state_key`).
|
||||
fn room_state_get(
|
||||
&self, room_id: &RoomId, event_type: &StateEventType, state_key: &str,
|
||||
) -> Result<Option<Arc<PduEvent>>> {
|
||||
if let Some(current_shortstatehash) = services().rooms.state.get_room_shortstatehash(room_id)? {
|
||||
self.state_get(current_shortstatehash, event_type, state_key)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -4,7 +4,7 @@ use std::{
|
|||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use lru_cache::LruCache;
|
||||
use ruma::{
|
||||
events::{
|
||||
|
|
|
@ -1,15 +1,21 @@
|
|||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use itertools::Itertools;
|
||||
use ruma::{
|
||||
events::{AnyStrippedStateEvent, AnySyncStateEvent},
|
||||
serde::Raw,
|
||||
OwnedRoomId, OwnedServerName, OwnedUserId, RoomId, ServerName, UserId,
|
||||
};
|
||||
use tracing::error;
|
||||
|
||||
use crate::{service::appservice::RegistrationInfo, Result};
|
||||
use crate::{
|
||||
appservice::RegistrationInfo,
|
||||
services, user_is_local,
|
||||
utils::{self},
|
||||
Error, KeyValueDatabase, Result,
|
||||
};
|
||||
|
||||
type StrippedStateEventIter<'a> = Box<dyn Iterator<Item = Result<(OwnedRoomId, Vec<Raw<AnyStrippedStateEvent>>)>> + 'a>;
|
||||
|
||||
type AnySyncStateEventIter<'a> = Box<dyn Iterator<Item = Result<(OwnedRoomId, Vec<Raw<AnySyncStateEvent>>)>> + 'a>;
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
|
@ -91,3 +97,609 @@ pub trait Data: Send + Sync {
|
|||
#[allow(dead_code)]
|
||||
fn add_servers_invite_via(&self, room_id: &RoomId, servers: &[OwnedServerName]) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn mark_as_once_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
self.roomuseroncejoinedids.insert(&userroom_id, &[])
|
||||
}
|
||||
|
||||
fn mark_as_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
|
||||
let roomid = room_id.as_bytes().to_vec();
|
||||
let mut roomid_prefix = room_id.as_bytes().to_vec();
|
||||
roomid_prefix.push(0xFF);
|
||||
|
||||
let mut roomuser_id = roomid_prefix.clone();
|
||||
roomuser_id.push(0xFF);
|
||||
roomuser_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_joined.insert(&userroom_id, &[])?;
|
||||
self.roomuserid_joined.insert(&roomuser_id, &[])?;
|
||||
self.userroomid_invitestate.remove(&userroom_id)?;
|
||||
self.roomuserid_invitecount.remove(&roomuser_id)?;
|
||||
self.userroomid_leftstate.remove(&userroom_id)?;
|
||||
self.roomuserid_leftcount.remove(&roomuser_id)?;
|
||||
|
||||
if self
|
||||
.roomuserid_joined
|
||||
.scan_prefix(roomid_prefix.clone())
|
||||
.count() == 0
|
||||
&& self
|
||||
.roomuserid_invitecount
|
||||
.scan_prefix(roomid_prefix)
|
||||
.count() == 0
|
||||
{
|
||||
self.roomid_inviteviaservers.remove(&roomid)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn mark_as_invited(
|
||||
&self, user_id: &UserId, room_id: &RoomId, last_state: Option<Vec<Raw<AnyStrippedStateEvent>>>,
|
||||
invite_via: Option<Vec<OwnedServerName>>,
|
||||
) -> Result<()> {
|
||||
let mut roomuser_id = room_id.as_bytes().to_vec();
|
||||
roomuser_id.push(0xFF);
|
||||
roomuser_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_invitestate.insert(
|
||||
&userroom_id,
|
||||
&serde_json::to_vec(&last_state.unwrap_or_default()).expect("state to bytes always works"),
|
||||
)?;
|
||||
self.roomuserid_invitecount
|
||||
.insert(&roomuser_id, &services().globals.next_count()?.to_be_bytes())?;
|
||||
self.userroomid_joined.remove(&userroom_id)?;
|
||||
self.roomuserid_joined.remove(&roomuser_id)?;
|
||||
self.userroomid_leftstate.remove(&userroom_id)?;
|
||||
self.roomuserid_leftcount.remove(&roomuser_id)?;
|
||||
|
||||
if let Some(servers) = invite_via {
|
||||
let mut prev_servers = self.servers_invite_via(room_id)?.unwrap_or(Vec::new());
|
||||
#[allow(clippy::redundant_clone)] // this is a necessary clone?
|
||||
prev_servers.append(servers.clone().as_mut());
|
||||
let servers = prev_servers.iter().rev().unique().rev().collect_vec();
|
||||
|
||||
let servers = servers
|
||||
.iter()
|
||||
.map(|server| server.as_bytes())
|
||||
.collect_vec()
|
||||
.join(&[0xFF][..]);
|
||||
|
||||
self.roomid_inviteviaservers
|
||||
.insert(room_id.as_bytes(), &servers)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn mark_as_left(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
|
||||
let roomid = room_id.as_bytes().to_vec();
|
||||
let mut roomid_prefix = room_id.as_bytes().to_vec();
|
||||
roomid_prefix.push(0xFF);
|
||||
|
||||
let mut roomuser_id = roomid_prefix.clone();
|
||||
roomuser_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_leftstate.insert(
|
||||
&userroom_id,
|
||||
&serde_json::to_vec(&Vec::<Raw<AnySyncStateEvent>>::new()).unwrap(),
|
||||
)?; // TODO
|
||||
self.roomuserid_leftcount
|
||||
.insert(&roomuser_id, &services().globals.next_count()?.to_be_bytes())?;
|
||||
self.userroomid_joined.remove(&userroom_id)?;
|
||||
self.roomuserid_joined.remove(&roomuser_id)?;
|
||||
self.userroomid_invitestate.remove(&userroom_id)?;
|
||||
self.roomuserid_invitecount.remove(&roomuser_id)?;
|
||||
|
||||
if self
|
||||
.roomuserid_joined
|
||||
.scan_prefix(roomid_prefix.clone())
|
||||
.count() == 0
|
||||
&& self
|
||||
.roomuserid_invitecount
|
||||
.scan_prefix(roomid_prefix)
|
||||
.count() == 0
|
||||
{
|
||||
self.roomid_inviteviaservers.remove(&roomid)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn update_joined_count(&self, room_id: &RoomId) -> Result<()> {
|
||||
let mut joinedcount = 0_u64;
|
||||
let mut invitedcount = 0_u64;
|
||||
let mut joined_servers = HashSet::new();
|
||||
let mut real_users = HashSet::new();
|
||||
|
||||
for joined in self.room_members(room_id).filter_map(Result::ok) {
|
||||
joined_servers.insert(joined.server_name().to_owned());
|
||||
if user_is_local(&joined) && !services().users.is_deactivated(&joined).unwrap_or(true) {
|
||||
real_users.insert(joined);
|
||||
}
|
||||
joinedcount = joinedcount.saturating_add(1);
|
||||
}
|
||||
|
||||
for _invited in self.room_members_invited(room_id).filter_map(Result::ok) {
|
||||
invitedcount = invitedcount.saturating_add(1);
|
||||
}
|
||||
|
||||
self.roomid_joinedcount
|
||||
.insert(room_id.as_bytes(), &joinedcount.to_be_bytes())?;
|
||||
|
||||
self.roomid_invitedcount
|
||||
.insert(room_id.as_bytes(), &invitedcount.to_be_bytes())?;
|
||||
|
||||
self.our_real_users_cache
|
||||
.write()
|
||||
.unwrap()
|
||||
.insert(room_id.to_owned(), Arc::new(real_users));
|
||||
|
||||
for old_joined_server in self.room_servers(room_id).filter_map(Result::ok) {
|
||||
if !joined_servers.remove(&old_joined_server) {
|
||||
// Server not in room anymore
|
||||
let mut roomserver_id = room_id.as_bytes().to_vec();
|
||||
roomserver_id.push(0xFF);
|
||||
roomserver_id.extend_from_slice(old_joined_server.as_bytes());
|
||||
|
||||
let mut serverroom_id = old_joined_server.as_bytes().to_vec();
|
||||
serverroom_id.push(0xFF);
|
||||
serverroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.roomserverids.remove(&roomserver_id)?;
|
||||
self.serverroomids.remove(&serverroom_id)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Now only new servers are in joined_servers anymore
|
||||
for server in joined_servers {
|
||||
let mut roomserver_id = room_id.as_bytes().to_vec();
|
||||
roomserver_id.push(0xFF);
|
||||
roomserver_id.extend_from_slice(server.as_bytes());
|
||||
|
||||
let mut serverroom_id = server.as_bytes().to_vec();
|
||||
serverroom_id.push(0xFF);
|
||||
serverroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.roomserverids.insert(&roomserver_id, &[])?;
|
||||
self.serverroomids.insert(&serverroom_id, &[])?;
|
||||
}
|
||||
|
||||
self.appservice_in_room_cache
|
||||
.write()
|
||||
.unwrap()
|
||||
.remove(room_id);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self, room_id))]
|
||||
fn get_our_real_users(&self, room_id: &RoomId) -> Result<Arc<HashSet<OwnedUserId>>> {
|
||||
let maybe = self
|
||||
.our_real_users_cache
|
||||
.read()
|
||||
.unwrap()
|
||||
.get(room_id)
|
||||
.cloned();
|
||||
if let Some(users) = maybe {
|
||||
Ok(users)
|
||||
} else {
|
||||
self.update_joined_count(room_id)?;
|
||||
Ok(Arc::clone(
|
||||
self.our_real_users_cache
|
||||
.read()
|
||||
.unwrap()
|
||||
.get(room_id)
|
||||
.unwrap(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self, room_id, appservice))]
|
||||
fn appservice_in_room(&self, room_id: &RoomId, appservice: &RegistrationInfo) -> Result<bool> {
|
||||
let maybe = self
|
||||
.appservice_in_room_cache
|
||||
.read()
|
||||
.unwrap()
|
||||
.get(room_id)
|
||||
.and_then(|map| map.get(&appservice.registration.id))
|
||||
.copied();
|
||||
|
||||
if let Some(b) = maybe {
|
||||
Ok(b)
|
||||
} else {
|
||||
let bridge_user_id = UserId::parse_with_server_name(
|
||||
appservice.registration.sender_localpart.as_str(),
|
||||
services().globals.server_name(),
|
||||
)
|
||||
.ok();
|
||||
|
||||
let in_room = bridge_user_id.map_or(false, |id| self.is_joined(&id, room_id).unwrap_or(false))
|
||||
|| self
|
||||
.room_members(room_id)
|
||||
.any(|userid| userid.map_or(false, |userid| appservice.users.is_match(userid.as_str())));
|
||||
|
||||
self.appservice_in_room_cache
|
||||
.write()
|
||||
.unwrap()
|
||||
.entry(room_id.to_owned())
|
||||
.or_default()
|
||||
.insert(appservice.registration.id.clone(), in_room);
|
||||
|
||||
Ok(in_room)
|
||||
}
|
||||
}
|
||||
|
||||
/// Makes a user forget a room.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn forget(&self, room_id: &RoomId, user_id: &UserId) -> Result<()> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
let mut roomuser_id = room_id.as_bytes().to_vec();
|
||||
roomuser_id.push(0xFF);
|
||||
roomuser_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.userroomid_leftstate.remove(&userroom_id)?;
|
||||
self.roomuserid_leftcount.remove(&roomuser_id)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns an iterator of all servers participating in this room.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_servers<'a>(&'a self, room_id: &RoomId) -> Box<dyn Iterator<Item = Result<OwnedServerName>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(self.roomserverids.scan_prefix(prefix).map(|(key, _)| {
|
||||
ServerName::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Server name in roomserverids is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Server name in roomserverids is invalid."))
|
||||
}))
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn server_in_room(&self, server: &ServerName, room_id: &RoomId) -> Result<bool> {
|
||||
let mut key = server.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.serverroomids.get(&key).map(|o| o.is_some())
|
||||
}
|
||||
|
||||
/// Returns an iterator of all rooms a server participates in (as far as we
|
||||
/// know).
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn server_rooms<'a>(&'a self, server: &ServerName) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a> {
|
||||
let mut prefix = server.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(self.serverroomids.scan_prefix(prefix).map(|(key, _)| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("RoomId in serverroomids is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("RoomId in serverroomids is invalid."))
|
||||
}))
|
||||
}
|
||||
|
||||
/// Returns an iterator over all joined members of a room.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_members<'a>(&'a self, room_id: &RoomId) -> Box<dyn Iterator<Item = Result<OwnedUserId>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(self.roomuserid_joined.scan_prefix(prefix).map(|(key, _)| {
|
||||
UserId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in roomuserid_joined is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in roomuserid_joined is invalid."))
|
||||
}))
|
||||
}
|
||||
|
||||
/// Returns the number of users which are currently in a room
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_joined_count(&self, room_id: &RoomId) -> Result<Option<u64>> {
|
||||
self.roomid_joinedcount
|
||||
.get(room_id.as_bytes())?
|
||||
.map(|b| utils::u64_from_bytes(&b).map_err(|_| Error::bad_database("Invalid joinedcount in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns the number of users which are currently invited to a room
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_invited_count(&self, room_id: &RoomId) -> Result<Option<u64>> {
|
||||
self.roomid_invitedcount
|
||||
.get(room_id.as_bytes())?
|
||||
.map(|b| utils::u64_from_bytes(&b).map_err(|_| Error::bad_database("Invalid joinedcount in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns an iterator over all User IDs who ever joined a room.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_useroncejoined<'a>(&'a self, room_id: &RoomId) -> Box<dyn Iterator<Item = Result<OwnedUserId>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(
|
||||
self.roomuseroncejoinedids
|
||||
.scan_prefix(prefix)
|
||||
.map(|(key, _)| {
|
||||
UserId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in room_useroncejoined is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in room_useroncejoined is invalid."))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns an iterator over all invited members of a room.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn room_members_invited<'a>(&'a self, room_id: &RoomId) -> Box<dyn Iterator<Item = Result<OwnedUserId>> + 'a> {
|
||||
let mut prefix = room_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(
|
||||
self.roomuserid_invitecount
|
||||
.scan_prefix(prefix)
|
||||
.map(|(key, _)| {
|
||||
UserId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in roomuserid_invited is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("User ID in roomuserid_invited is invalid."))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn get_invite_count(&self, room_id: &RoomId, user_id: &UserId) -> Result<Option<u64>> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.roomuserid_invitecount
|
||||
.get(&key)?
|
||||
.map_or(Ok(None), |bytes| {
|
||||
Ok(Some(
|
||||
utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid invitecount in db."))?,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn get_left_count(&self, room_id: &RoomId, user_id: &UserId) -> Result<Option<u64>> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.roomuserid_leftcount
|
||||
.get(&key)?
|
||||
.map(|bytes| utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid leftcount in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns an iterator over all rooms this user joined.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn rooms_joined<'a>(&'a self, user_id: &UserId) -> Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a> {
|
||||
Box::new(
|
||||
self.userroomid_joined
|
||||
.scan_prefix(user_id.as_bytes().to_vec())
|
||||
.map(|(key, _)| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_joined is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_joined is invalid."))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns an iterator over all rooms a user was invited to.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn rooms_invited<'a>(&'a self, user_id: &UserId) -> StrippedStateEventIter<'a> {
|
||||
let mut prefix = user_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(
|
||||
self.userroomid_invitestate
|
||||
.scan_prefix(prefix)
|
||||
.map(|(key, state)| {
|
||||
let room_id = RoomId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?;
|
||||
|
||||
let state = serde_json::from_slice(&state)
|
||||
.map_err(|_| Error::bad_database("Invalid state in userroomid_invitestate."))?;
|
||||
|
||||
Ok((room_id, state))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn invite_state(&self, user_id: &UserId, room_id: &RoomId) -> Result<Option<Vec<Raw<AnyStrippedStateEvent>>>> {
|
||||
let mut key = user_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_invitestate
|
||||
.get(&key)?
|
||||
.map(|state| {
|
||||
let state = serde_json::from_slice(&state)
|
||||
.map_err(|_| Error::bad_database("Invalid state in userroomid_invitestate."))?;
|
||||
|
||||
Ok(state)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn left_state(&self, user_id: &UserId, room_id: &RoomId) -> Result<Option<Vec<Raw<AnyStrippedStateEvent>>>> {
|
||||
let mut key = user_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_leftstate
|
||||
.get(&key)?
|
||||
.map(|state| {
|
||||
let state = serde_json::from_slice(&state)
|
||||
.map_err(|_| Error::bad_database("Invalid state in userroomid_leftstate."))?;
|
||||
|
||||
Ok(state)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns an iterator over all rooms a user left.
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn rooms_left<'a>(&'a self, user_id: &UserId) -> AnySyncStateEventIter<'a> {
|
||||
let mut prefix = user_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
Box::new(
|
||||
self.userroomid_leftstate
|
||||
.scan_prefix(prefix)
|
||||
.map(|(key, state)| {
|
||||
let room_id = RoomId::parse(
|
||||
utils::string_from_bytes(
|
||||
key.rsplit(|&b| b == 0xFF)
|
||||
.next()
|
||||
.expect("rsplit always returns an element"),
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid unicode."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Room ID in userroomid_invited is invalid."))?;
|
||||
|
||||
let state = serde_json::from_slice(&state)
|
||||
.map_err(|_| Error::bad_database("Invalid state in userroomid_leftstate."))?;
|
||||
|
||||
Ok((room_id, state))
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn once_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result<bool> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
Ok(self.roomuseroncejoinedids.get(&userroom_id)?.is_some())
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn is_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result<bool> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
Ok(self.userroomid_joined.get(&userroom_id)?.is_some())
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn is_invited(&self, user_id: &UserId, room_id: &RoomId) -> Result<bool> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
Ok(self.userroomid_invitestate.get(&userroom_id)?.is_some())
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn is_left(&self, user_id: &UserId, room_id: &RoomId) -> Result<bool> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
Ok(self.userroomid_leftstate.get(&userroom_id)?.is_some())
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn servers_invite_via(&self, room_id: &RoomId) -> Result<Option<Vec<OwnedServerName>>> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
|
||||
self.roomid_inviteviaservers
|
||||
.get(&key)?
|
||||
.map(|servers| {
|
||||
let state = serde_json::from_slice(&servers).map_err(|e| {
|
||||
error!("Invalid state in userroomid_leftstate: {e}");
|
||||
Error::bad_database("Invalid state in userroomid_leftstate.")
|
||||
})?;
|
||||
|
||||
Ok(state)
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip(self))]
|
||||
fn add_servers_invite_via(&self, room_id: &RoomId, servers: &[OwnedServerName]) -> Result<()> {
|
||||
let mut prev_servers = self.servers_invite_via(room_id)?.unwrap_or(Vec::new());
|
||||
prev_servers.append(servers.to_owned().as_mut());
|
||||
|
||||
let servers = prev_servers.iter().rev().unique().rev().collect_vec();
|
||||
|
||||
let servers = servers
|
||||
.iter()
|
||||
.map(|server| server.as_bytes())
|
||||
.collect_vec()
|
||||
.join(&[0xFF][..]);
|
||||
|
||||
self.roomid_inviteviaservers
|
||||
.insert(room_id.as_bytes(), &servers)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use itertools::Itertools;
|
||||
use ruma::{
|
||||
events::{
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
use std::{collections::HashSet, sync::Arc};
|
||||
use std::{collections::HashSet, mem::size_of, sync::Arc};
|
||||
|
||||
use super::CompressedStateEvent;
|
||||
use crate::Result;
|
||||
use crate::{utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub struct StateDiff {
|
||||
pub parent: Option<u64>,
|
||||
|
@ -13,3 +13,60 @@ pub trait Data: Send + Sync {
|
|||
fn get_statediff(&self, shortstatehash: u64) -> Result<StateDiff>;
|
||||
fn save_statediff(&self, shortstatehash: u64, diff: StateDiff) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn get_statediff(&self, shortstatehash: u64) -> Result<StateDiff> {
|
||||
let value = self
|
||||
.shortstatehash_statediff
|
||||
.get(&shortstatehash.to_be_bytes())?
|
||||
.ok_or_else(|| Error::bad_database("State hash does not exist"))?;
|
||||
let parent = utils::u64_from_bytes(&value[0..size_of::<u64>()]).expect("bytes have right length");
|
||||
let parent = if parent != 0 {
|
||||
Some(parent)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut add_mode = true;
|
||||
let mut added = HashSet::new();
|
||||
let mut removed = HashSet::new();
|
||||
|
||||
let mut i = size_of::<u64>();
|
||||
while let Some(v) = value.get(i..i + 2 * size_of::<u64>()) {
|
||||
if add_mode && v.starts_with(&0_u64.to_be_bytes()) {
|
||||
add_mode = false;
|
||||
i += size_of::<u64>();
|
||||
continue;
|
||||
}
|
||||
if add_mode {
|
||||
added.insert(v.try_into().expect("we checked the size above"));
|
||||
} else {
|
||||
removed.insert(v.try_into().expect("we checked the size above"));
|
||||
}
|
||||
i += 2 * size_of::<u64>();
|
||||
}
|
||||
|
||||
Ok(StateDiff {
|
||||
parent,
|
||||
added: Arc::new(added),
|
||||
removed: Arc::new(removed),
|
||||
})
|
||||
}
|
||||
|
||||
fn save_statediff(&self, shortstatehash: u64, diff: StateDiff) -> Result<()> {
|
||||
let mut value = diff.parent.unwrap_or(0).to_be_bytes().to_vec();
|
||||
for new in diff.added.iter() {
|
||||
value.extend_from_slice(&new[..]);
|
||||
}
|
||||
|
||||
if !diff.removed.is_empty() {
|
||||
value.extend_from_slice(&0_u64.to_be_bytes());
|
||||
for removed in diff.removed.iter() {
|
||||
value.extend_from_slice(&removed[..]);
|
||||
}
|
||||
}
|
||||
|
||||
self.shortstatehash_statediff
|
||||
.insert(&shortstatehash.to_be_bytes(), &value)
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,11 +1,11 @@
|
|||
pub mod data;
|
||||
mod data;
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
mem::size_of,
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use lru_cache::LruCache;
|
||||
use ruma::{EventId, RoomId};
|
||||
|
||||
|
|
|
@ -1,6 +1,8 @@
|
|||
use std::mem;
|
||||
|
||||
use ruma::{api::client::threads::get_threads::v1::IncludeThreads, OwnedUserId, RoomId, UserId};
|
||||
|
||||
use crate::{PduEvent, Result};
|
||||
use crate::{services, utils, Error, KeyValueDatabase, PduEvent, Result};
|
||||
|
||||
type PduEventIterResult<'a> = Result<Box<dyn Iterator<Item = Result<(u64, PduEvent)>> + 'a>>;
|
||||
|
||||
|
@ -12,3 +14,71 @@ pub trait Data: Send + Sync {
|
|||
fn update_participants(&self, root_id: &[u8], participants: &[OwnedUserId]) -> Result<()>;
|
||||
fn get_participants(&self, root_id: &[u8]) -> Result<Option<Vec<OwnedUserId>>>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn threads_until<'a>(
|
||||
&'a self, user_id: &'a UserId, room_id: &'a RoomId, until: u64, _include: &'a IncludeThreads,
|
||||
) -> PduEventIterResult<'a> {
|
||||
let prefix = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortroomid(room_id)?
|
||||
.expect("room exists")
|
||||
.to_be_bytes()
|
||||
.to_vec();
|
||||
|
||||
let mut current = prefix.clone();
|
||||
current.extend_from_slice(&(until - 1).to_be_bytes());
|
||||
|
||||
Ok(Box::new(
|
||||
self.threadid_userids
|
||||
.iter_from(¤t, true)
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix))
|
||||
.map(move |(pduid, _users)| {
|
||||
let count = utils::u64_from_bytes(&pduid[(mem::size_of::<u64>())..])
|
||||
.map_err(|_| Error::bad_database("Invalid pduid in threadid_userids."))?;
|
||||
let mut pdu = services()
|
||||
.rooms
|
||||
.timeline
|
||||
.get_pdu_from_id(&pduid)?
|
||||
.ok_or_else(|| Error::bad_database("Invalid pduid reference in threadid_userids"))?;
|
||||
if pdu.sender != user_id {
|
||||
pdu.remove_transaction_id()?;
|
||||
}
|
||||
Ok((count, pdu))
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn update_participants(&self, root_id: &[u8], participants: &[OwnedUserId]) -> Result<()> {
|
||||
let users = participants
|
||||
.iter()
|
||||
.map(|user| user.as_bytes())
|
||||
.collect::<Vec<_>>()
|
||||
.join(&[0xFF][..]);
|
||||
|
||||
self.threadid_userids.insert(root_id, &users)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_participants(&self, root_id: &[u8]) -> Result<Option<Vec<OwnedUserId>>> {
|
||||
if let Some(users) = self.threadid_userids.get(root_id)? {
|
||||
Ok(Some(
|
||||
users
|
||||
.split(|b| *b == 0xFF)
|
||||
.map(|bytes| {
|
||||
UserId::parse(
|
||||
utils::string_from_bytes(bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid UserId bytes in threadid_userids."))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Invalid UserId in threadid_userids."))
|
||||
})
|
||||
.filter_map(Result::ok)
|
||||
.collect(),
|
||||
))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{
|
||||
api::client::{error::ErrorKind, threads::get_threads::v1::IncludeThreads},
|
||||
events::relation::BundledThread,
|
||||
|
|
|
@ -1,9 +1,10 @@
|
|||
use std::sync::Arc;
|
||||
use std::{collections::hash_map, mem::size_of, sync::Arc};
|
||||
|
||||
use ruma::{CanonicalJsonObject, EventId, OwnedUserId, RoomId, UserId};
|
||||
use ruma::{api::client::error::ErrorKind, CanonicalJsonObject, EventId, OwnedUserId, RoomId, UserId};
|
||||
use tracing::error;
|
||||
|
||||
use super::PduCount;
|
||||
use crate::{PduEvent, Result};
|
||||
use crate::{services, utils, Error, KeyValueDatabase, PduEvent, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn last_timeline_count(&self, sender_user: &UserId, room_id: &RoomId) -> Result<PduCount>;
|
||||
|
@ -66,3 +67,292 @@ pub trait Data: Send + Sync {
|
|||
&self, room_id: &RoomId, notifies: Vec<OwnedUserId>, highlights: Vec<OwnedUserId>,
|
||||
) -> Result<()>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn last_timeline_count(&self, sender_user: &UserId, room_id: &RoomId) -> Result<PduCount> {
|
||||
match self
|
||||
.lasttimelinecount_cache
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(room_id.to_owned())
|
||||
{
|
||||
hash_map::Entry::Vacant(v) => {
|
||||
if let Some(last_count) = self
|
||||
.pdus_until(sender_user, room_id, PduCount::max())?
|
||||
.find_map(|r| {
|
||||
// Filter out buggy events
|
||||
if r.is_err() {
|
||||
error!("Bad pdu in pdus_since: {:?}", r);
|
||||
}
|
||||
r.ok()
|
||||
}) {
|
||||
Ok(*v.insert(last_count.0))
|
||||
} else {
|
||||
Ok(PduCount::Normal(0))
|
||||
}
|
||||
},
|
||||
hash_map::Entry::Occupied(o) => Ok(*o.get()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the `count` of this pdu's id.
|
||||
fn get_pdu_count(&self, event_id: &EventId) -> Result<Option<PduCount>> {
|
||||
self.eventid_pduid
|
||||
.get(event_id.as_bytes())?
|
||||
.map(|pdu_id| pdu_count(&pdu_id))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns the json of a pdu.
|
||||
fn get_pdu_json(&self, event_id: &EventId) -> Result<Option<CanonicalJsonObject>> {
|
||||
self.get_non_outlier_pdu_json(event_id)?.map_or_else(
|
||||
|| {
|
||||
self.eventid_outlierpdu
|
||||
.get(event_id.as_bytes())?
|
||||
.map(|pdu| serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db.")))
|
||||
.transpose()
|
||||
},
|
||||
|x| Ok(Some(x)),
|
||||
)
|
||||
}
|
||||
|
||||
/// Returns the json of a pdu.
|
||||
fn get_non_outlier_pdu_json(&self, event_id: &EventId) -> Result<Option<CanonicalJsonObject>> {
|
||||
self.eventid_pduid
|
||||
.get(event_id.as_bytes())?
|
||||
.map(|pduid| {
|
||||
self.pduid_pdu
|
||||
.get(&pduid)?
|
||||
.ok_or_else(|| Error::bad_database("Invalid pduid in eventid_pduid."))
|
||||
})
|
||||
.transpose()?
|
||||
.map(|pdu| serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns the pdu's id.
|
||||
fn get_pdu_id(&self, event_id: &EventId) -> Result<Option<Vec<u8>>> { self.eventid_pduid.get(event_id.as_bytes()) }
|
||||
|
||||
/// Returns the pdu.
|
||||
fn get_non_outlier_pdu(&self, event_id: &EventId) -> Result<Option<PduEvent>> {
|
||||
self.eventid_pduid
|
||||
.get(event_id.as_bytes())?
|
||||
.map(|pduid| {
|
||||
self.pduid_pdu
|
||||
.get(&pduid)?
|
||||
.ok_or_else(|| Error::bad_database("Invalid pduid in eventid_pduid."))
|
||||
})
|
||||
.transpose()?
|
||||
.map(|pdu| serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db.")))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
/// Returns the pdu.
|
||||
///
|
||||
/// Checks the `eventid_outlierpdu` Tree if not found in the timeline.
|
||||
fn get_pdu(&self, event_id: &EventId) -> Result<Option<Arc<PduEvent>>> {
|
||||
if let Some(pdu) = self
|
||||
.get_non_outlier_pdu(event_id)?
|
||||
.map_or_else(
|
||||
|| {
|
||||
self.eventid_outlierpdu
|
||||
.get(event_id.as_bytes())?
|
||||
.map(|pdu| serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db.")))
|
||||
.transpose()
|
||||
},
|
||||
|x| Ok(Some(x)),
|
||||
)?
|
||||
.map(Arc::new)
|
||||
{
|
||||
Ok(Some(pdu))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the pdu.
|
||||
///
|
||||
/// This does __NOT__ check the outliers `Tree`.
|
||||
fn get_pdu_from_id(&self, pdu_id: &[u8]) -> Result<Option<PduEvent>> {
|
||||
self.pduid_pdu.get(pdu_id)?.map_or(Ok(None), |pdu| {
|
||||
Ok(Some(
|
||||
serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db."))?,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the pdu as a `BTreeMap<String, CanonicalJsonValue>`.
|
||||
fn get_pdu_json_from_id(&self, pdu_id: &[u8]) -> Result<Option<CanonicalJsonObject>> {
|
||||
self.pduid_pdu.get(pdu_id)?.map_or(Ok(None), |pdu| {
|
||||
Ok(Some(
|
||||
serde_json::from_slice(&pdu).map_err(|_| Error::bad_database("Invalid PDU in db."))?,
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn append_pdu(&self, pdu_id: &[u8], pdu: &PduEvent, json: &CanonicalJsonObject, count: u64) -> Result<()> {
|
||||
self.pduid_pdu.insert(
|
||||
pdu_id,
|
||||
&serde_json::to_vec(json).expect("CanonicalJsonObject is always a valid"),
|
||||
)?;
|
||||
|
||||
self.lasttimelinecount_cache
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(pdu.room_id.clone(), PduCount::Normal(count));
|
||||
|
||||
self.eventid_pduid.insert(pdu.event_id.as_bytes(), pdu_id)?;
|
||||
self.eventid_outlierpdu.remove(pdu.event_id.as_bytes())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn prepend_backfill_pdu(&self, pdu_id: &[u8], event_id: &EventId, json: &CanonicalJsonObject) -> Result<()> {
|
||||
self.pduid_pdu.insert(
|
||||
pdu_id,
|
||||
&serde_json::to_vec(json).expect("CanonicalJsonObject is always a valid"),
|
||||
)?;
|
||||
|
||||
self.eventid_pduid.insert(event_id.as_bytes(), pdu_id)?;
|
||||
self.eventid_outlierpdu.remove(event_id.as_bytes())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Removes a pdu and creates a new one with the same id.
|
||||
fn replace_pdu(&self, pdu_id: &[u8], pdu_json: &CanonicalJsonObject, _pdu: &PduEvent) -> Result<()> {
|
||||
if self.pduid_pdu.get(pdu_id)?.is_some() {
|
||||
self.pduid_pdu.insert(
|
||||
pdu_id,
|
||||
&serde_json::to_vec(pdu_json).expect("CanonicalJsonObject is always a valid"),
|
||||
)?;
|
||||
} else {
|
||||
return Err(Error::BadRequest(ErrorKind::NotFound, "PDU does not exist."));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns an iterator over all events and their tokens in a room that
|
||||
/// happened before the event with id `until` in reverse-chronological
|
||||
/// order.
|
||||
fn pdus_until<'a>(
|
||||
&'a self, user_id: &UserId, room_id: &RoomId, until: PduCount,
|
||||
) -> Result<Box<dyn Iterator<Item = Result<(PduCount, PduEvent)>> + 'a>> {
|
||||
let (prefix, current) = count_to_id(room_id, until, 1, true)?;
|
||||
|
||||
let user_id = user_id.to_owned();
|
||||
|
||||
Ok(Box::new(
|
||||
self.pduid_pdu
|
||||
.iter_from(¤t, true)
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix))
|
||||
.map(move |(pdu_id, v)| {
|
||||
let mut pdu = serde_json::from_slice::<PduEvent>(&v)
|
||||
.map_err(|_| Error::bad_database("PDU in db is invalid."))?;
|
||||
if pdu.sender != user_id {
|
||||
pdu.remove_transaction_id()?;
|
||||
}
|
||||
pdu.add_age()?;
|
||||
let count = pdu_count(&pdu_id)?;
|
||||
Ok((count, pdu))
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn pdus_after<'a>(
|
||||
&'a self, user_id: &UserId, room_id: &RoomId, from: PduCount,
|
||||
) -> Result<Box<dyn Iterator<Item = Result<(PduCount, PduEvent)>> + 'a>> {
|
||||
let (prefix, current) = count_to_id(room_id, from, 1, false)?;
|
||||
|
||||
let user_id = user_id.to_owned();
|
||||
|
||||
Ok(Box::new(
|
||||
self.pduid_pdu
|
||||
.iter_from(¤t, false)
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix))
|
||||
.map(move |(pdu_id, v)| {
|
||||
let mut pdu = serde_json::from_slice::<PduEvent>(&v)
|
||||
.map_err(|_| Error::bad_database("PDU in db is invalid."))?;
|
||||
if pdu.sender != user_id {
|
||||
pdu.remove_transaction_id()?;
|
||||
}
|
||||
pdu.add_age()?;
|
||||
let count = pdu_count(&pdu_id)?;
|
||||
Ok((count, pdu))
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn increment_notification_counts(
|
||||
&self, room_id: &RoomId, notifies: Vec<OwnedUserId>, highlights: Vec<OwnedUserId>,
|
||||
) -> Result<()> {
|
||||
let mut notifies_batch = Vec::new();
|
||||
let mut highlights_batch = Vec::new();
|
||||
for user in notifies {
|
||||
let mut userroom_id = user.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
notifies_batch.push(userroom_id);
|
||||
}
|
||||
for user in highlights {
|
||||
let mut userroom_id = user.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
highlights_batch.push(userroom_id);
|
||||
}
|
||||
|
||||
self.userroomid_notificationcount
|
||||
.increment_batch(&mut notifies_batch.into_iter())?;
|
||||
self.userroomid_highlightcount
|
||||
.increment_batch(&mut highlights_batch.into_iter())?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the `count` of this pdu's id.
|
||||
fn pdu_count(pdu_id: &[u8]) -> Result<PduCount> {
|
||||
let last_u64 = utils::u64_from_bytes(&pdu_id[pdu_id.len() - size_of::<u64>()..])
|
||||
.map_err(|_| Error::bad_database("PDU has invalid count bytes."))?;
|
||||
let second_last_u64 =
|
||||
utils::u64_from_bytes(&pdu_id[pdu_id.len() - 2 * size_of::<u64>()..pdu_id.len() - size_of::<u64>()]);
|
||||
|
||||
if matches!(second_last_u64, Ok(0)) {
|
||||
Ok(PduCount::Backfilled(u64::MAX.saturating_sub(last_u64)))
|
||||
} else {
|
||||
Ok(PduCount::Normal(last_u64))
|
||||
}
|
||||
}
|
||||
|
||||
fn count_to_id(room_id: &RoomId, count: PduCount, offset: u64, subtract: bool) -> Result<(Vec<u8>, Vec<u8>)> {
|
||||
let prefix = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortroomid(room_id)?
|
||||
.ok_or_else(|| Error::bad_database("Looked for bad shortroomid in timeline"))?
|
||||
.to_be_bytes()
|
||||
.to_vec();
|
||||
let mut pdu_id = prefix.clone();
|
||||
// +1 so we don't send the base event
|
||||
let count_raw = match count {
|
||||
PduCount::Normal(x) => {
|
||||
if subtract {
|
||||
x.saturating_sub(offset)
|
||||
} else {
|
||||
x.saturating_add(offset)
|
||||
}
|
||||
},
|
||||
PduCount::Backfilled(x) => {
|
||||
pdu_id.extend_from_slice(&0_u64.to_be_bytes());
|
||||
let num = u64::MAX.saturating_sub(x);
|
||||
if subtract {
|
||||
num.saturating_sub(offset)
|
||||
} else {
|
||||
num.saturating_add(offset)
|
||||
}
|
||||
},
|
||||
};
|
||||
pdu_id.extend_from_slice(&count_raw.to_be_bytes());
|
||||
|
||||
Ok((prefix, pdu_id))
|
||||
}
|
||||
|
|
|
@ -1,11 +1,11 @@
|
|||
pub mod data;
|
||||
mod data;
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, HashMap, HashSet},
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use rand::prelude::SliceRandom;
|
||||
use ruma::{
|
||||
api::{client::error::ErrorKind, federation},
|
||||
|
@ -195,7 +195,7 @@ impl Service {
|
|||
state_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room state mutex
|
||||
) -> Result<Vec<u8>> {
|
||||
// Coalesce database writes for the remainder of this scope.
|
||||
let _cork = services().globals.db.cork_and_flush()?;
|
||||
let _cork = services().globals.cork_and_flush()?;
|
||||
|
||||
let shortroomid = services()
|
||||
.rooms
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use ruma::{OwnedRoomId, OwnedUserId, RoomId, UserId};
|
||||
|
||||
use crate::Result;
|
||||
use crate::{services, utils, Error, KeyValueDatabase, Result};
|
||||
|
||||
pub trait Data: Send + Sync {
|
||||
fn reset_notification_counts(&self, user_id: &UserId, room_id: &RoomId) -> Result<()>;
|
||||
|
@ -20,3 +20,137 @@ pub trait Data: Send + Sync {
|
|||
&'a self, users: Vec<OwnedUserId>,
|
||||
) -> Result<Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a>>;
|
||||
}
|
||||
|
||||
impl Data for KeyValueDatabase {
|
||||
fn reset_notification_counts(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
let mut roomuser_id = room_id.as_bytes().to_vec();
|
||||
roomuser_id.push(0xFF);
|
||||
roomuser_id.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
self.userroomid_notificationcount
|
||||
.insert(&userroom_id, &0_u64.to_be_bytes())?;
|
||||
self.userroomid_highlightcount
|
||||
.insert(&userroom_id, &0_u64.to_be_bytes())?;
|
||||
|
||||
self.roomuserid_lastnotificationread
|
||||
.insert(&roomuser_id, &services().globals.next_count()?.to_be_bytes())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn notification_count(&self, user_id: &UserId, room_id: &RoomId) -> Result<u64> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_notificationcount
|
||||
.get(&userroom_id)?
|
||||
.map_or(Ok(0), |bytes| {
|
||||
utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid notification count in db."))
|
||||
})
|
||||
}
|
||||
|
||||
fn highlight_count(&self, user_id: &UserId, room_id: &RoomId) -> Result<u64> {
|
||||
let mut userroom_id = user_id.as_bytes().to_vec();
|
||||
userroom_id.push(0xFF);
|
||||
userroom_id.extend_from_slice(room_id.as_bytes());
|
||||
|
||||
self.userroomid_highlightcount
|
||||
.get(&userroom_id)?
|
||||
.map_or(Ok(0), |bytes| {
|
||||
utils::u64_from_bytes(&bytes).map_err(|_| Error::bad_database("Invalid highlight count in db."))
|
||||
})
|
||||
}
|
||||
|
||||
fn last_notification_read(&self, user_id: &UserId, room_id: &RoomId) -> Result<u64> {
|
||||
let mut key = room_id.as_bytes().to_vec();
|
||||
key.push(0xFF);
|
||||
key.extend_from_slice(user_id.as_bytes());
|
||||
|
||||
Ok(self
|
||||
.roomuserid_lastnotificationread
|
||||
.get(&key)?
|
||||
.map(|bytes| {
|
||||
utils::u64_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Count in roomuserid_lastprivatereadupdate is invalid."))
|
||||
})
|
||||
.transpose()?
|
||||
.unwrap_or(0))
|
||||
}
|
||||
|
||||
fn associate_token_shortstatehash(&self, room_id: &RoomId, token: u64, shortstatehash: u64) -> Result<()> {
|
||||
let shortroomid = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortroomid(room_id)?
|
||||
.expect("room exists");
|
||||
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(&token.to_be_bytes());
|
||||
|
||||
self.roomsynctoken_shortstatehash
|
||||
.insert(&key, &shortstatehash.to_be_bytes())
|
||||
}
|
||||
|
||||
fn get_token_shortstatehash(&self, room_id: &RoomId, token: u64) -> Result<Option<u64>> {
|
||||
let shortroomid = services()
|
||||
.rooms
|
||||
.short
|
||||
.get_shortroomid(room_id)?
|
||||
.expect("room exists");
|
||||
|
||||
let mut key = shortroomid.to_be_bytes().to_vec();
|
||||
key.extend_from_slice(&token.to_be_bytes());
|
||||
|
||||
self.roomsynctoken_shortstatehash
|
||||
.get(&key)?
|
||||
.map(|bytes| {
|
||||
utils::u64_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid shortstatehash in roomsynctoken_shortstatehash"))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn get_shared_rooms<'a>(
|
||||
&'a self, users: Vec<OwnedUserId>,
|
||||
) -> Result<Box<dyn Iterator<Item = Result<OwnedRoomId>> + 'a>> {
|
||||
let iterators = users.into_iter().map(move |user_id| {
|
||||
let mut prefix = user_id.as_bytes().to_vec();
|
||||
prefix.push(0xFF);
|
||||
|
||||
self.userroomid_joined
|
||||
.scan_prefix(prefix)
|
||||
.map(|(key, _)| {
|
||||
let roomid_index = key
|
||||
.iter()
|
||||
.enumerate()
|
||||
.find(|(_, &b)| b == 0xFF)
|
||||
.ok_or_else(|| Error::bad_database("Invalid userroomid_joined in db."))?
|
||||
.0
|
||||
.saturating_add(1); // +1 because the room id starts AFTER the separator
|
||||
|
||||
let room_id = key[roomid_index..].to_vec();
|
||||
|
||||
Ok::<_, Error>(room_id)
|
||||
})
|
||||
.filter_map(Result::ok)
|
||||
});
|
||||
|
||||
// We use the default compare function because keys are sorted correctly (not
|
||||
// reversed)
|
||||
Ok(Box::new(
|
||||
utils::common_elements(iterators, Ord::cmp)
|
||||
.expect("users is not empty")
|
||||
.map(|bytes| {
|
||||
RoomId::parse(
|
||||
utils::string_from_bytes(&bytes)
|
||||
.map_err(|_| Error::bad_database("Invalid RoomId bytes in userroomid_joined"))?,
|
||||
)
|
||||
.map_err(|_| Error::bad_database("Invalid RoomId in userroomid_joined."))
|
||||
}),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
|
|
@ -2,7 +2,7 @@ mod data;
|
|||
|
||||
use std::sync::Arc;
|
||||
|
||||
pub use data::Data;
|
||||
use data::Data;
|
||||
use ruma::{OwnedRoomId, OwnedUserId, RoomId, UserId};
|
||||
|
||||
use crate::Result;
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue