use ruma::{ServerName, UserId}; use crate::{ database::KeyValueDatabase, service::{ self, sending::{OutgoingKind, SendingEventType}, }, services, utils, Error, Result, }; impl service::sending::Data for KeyValueDatabase { fn active_requests<'a>( &'a self, ) -> Box, OutgoingKind, SendingEventType)>> + 'a> { Box::new( self.servercurrentevent_data .iter() .map(|(key, v)| parse_servercurrentevent(&key, v).map(|(k, e)| (key, k, e))), ) } fn active_requests_for<'a>( &'a self, outgoing_kind: &OutgoingKind, ) -> Box, SendingEventType)>> + 'a> { let prefix = outgoing_kind.get_prefix(); Box::new( self.servercurrentevent_data .scan_prefix(prefix) .map(|(key, v)| parse_servercurrentevent(&key, v).map(|(_, e)| (key, e))), ) } fn delete_active_request(&self, key: Vec) -> Result<()> { self.servercurrentevent_data.remove(&key) } fn delete_all_active_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> { let prefix = outgoing_kind.get_prefix(); for (key, _) in self.servercurrentevent_data.scan_prefix(prefix) { self.servercurrentevent_data.remove(&key)?; } Ok(()) } fn delete_all_requests_for(&self, outgoing_kind: &OutgoingKind) -> Result<()> { let prefix = outgoing_kind.get_prefix(); for (key, _) in self.servercurrentevent_data.scan_prefix(prefix.clone()) { self.servercurrentevent_data.remove(&key).unwrap(); } for (key, _) in self.servernameevent_data.scan_prefix(prefix) { self.servernameevent_data.remove(&key).unwrap(); } Ok(()) } fn queue_requests( &self, requests: &[(&OutgoingKind, SendingEventType)], ) -> Result>> { let mut batch = Vec::new(); let mut keys = Vec::new(); for (outgoing_kind, event) in requests { let mut key = outgoing_kind.get_prefix(); if let SendingEventType::Pdu(value) = &event { key.extend_from_slice(value) } else { key.extend_from_slice(&services().globals.next_count()?.to_be_bytes()) } let value = if let SendingEventType::Edu(value) = &event { &**value } else { &[] }; batch.push((key.clone(), value.to_owned())); keys.push(key); } self.servernameevent_data .insert_batch(&mut batch.into_iter())?; Ok(keys) } fn queued_requests<'a>( &'a self, outgoing_kind: &OutgoingKind, ) -> Box)>> + 'a> { let prefix = outgoing_kind.get_prefix(); return Box::new( self.servernameevent_data .scan_prefix(prefix) .map(|(k, v)| parse_servercurrentevent(&k, v).map(|(_, ev)| (ev, k))), ); } fn mark_as_active(&self, events: &[(SendingEventType, Vec)]) -> Result<()> { for (e, key) in events { let value = if let SendingEventType::Edu(value) = &e { &**value } else { &[] }; self.servercurrentevent_data.insert(key, value)?; self.servernameevent_data.remove(key)?; } Ok(()) } fn set_latest_educount(&self, server_name: &ServerName, last_count: u64) -> Result<()> { self.servername_educount .insert(server_name.as_bytes(), &last_count.to_be_bytes()) } fn get_latest_educount(&self, server_name: &ServerName) -> Result { self.servername_educount .get(server_name.as_bytes())? .map_or(Ok(0), |bytes| { utils::u64_from_bytes(&bytes) .map_err(|_| Error::bad_database("Invalid u64 in servername_educount.")) }) } } #[tracing::instrument(skip(key))] fn parse_servercurrentevent( key: &[u8], value: Vec, ) -> Result<(OutgoingKind, SendingEventType)> { // Appservices start with a plus Ok::<_, Error>(if key.starts_with(b"+") { let mut parts = key[1..].splitn(2, |&b| b == 0xff); let server = parts.next().expect("splitn always returns one element"); let event = parts .next() .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?; let server = utils::string_from_bytes(server).map_err(|_| { Error::bad_database("Invalid server bytes in server_currenttransaction") })?; ( OutgoingKind::Appservice(server), if value.is_empty() { SendingEventType::Pdu(event.to_vec()) } else { SendingEventType::Edu(value) }, ) } else if key.starts_with(b"$") { let mut parts = key[1..].splitn(3, |&b| b == 0xff); let user = parts.next().expect("splitn always returns one element"); let user_string = utils::string_from_bytes(user) .map_err(|_| Error::bad_database("Invalid user string in servercurrentevent"))?; let user_id = UserId::parse(user_string) .map_err(|_| Error::bad_database("Invalid user id in servercurrentevent"))?; let pushkey = parts .next() .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?; let pushkey_string = utils::string_from_bytes(pushkey) .map_err(|_| Error::bad_database("Invalid pushkey in servercurrentevent"))?; let event = parts .next() .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?; ( OutgoingKind::Push(user_id, pushkey_string), if value.is_empty() { SendingEventType::Pdu(event.to_vec()) } else { // I'm pretty sure this should never be called SendingEventType::Edu(value) }, ) } else { let mut parts = key.splitn(2, |&b| b == 0xff); let server = parts.next().expect("splitn always returns one element"); let event = parts .next() .ok_or_else(|| Error::bad_database("Invalid bytes in servercurrentpdus."))?; let server = utils::string_from_bytes(server).map_err(|_| { Error::bad_database("Invalid server bytes in server_currenttransaction") })?; ( OutgoingKind::Normal(ServerName::parse(server).map_err(|_| { Error::bad_database("Invalid server string in server_currenttransaction") })?), if value.is_empty() { SendingEventType::Pdu(event.to_vec()) } else { SendingEventType::Edu(value) }, ) }) }