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< dyn Iterator, 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 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) }, ) }) }