continuwuity/src/database/key_value/sending.rs

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<Item = Result<(Vec<u8>, 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<dyn Iterator<Item = Result<(Vec<u8>, 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<u8>) -> 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<Vec<Vec<u8>>> {
		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<dyn Iterator<Item = Result<(SendingEventType, Vec<u8>)>> + '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<u8>)]) -> Result<()> {
		for (e, key) in events {
			if key.is_empty() {
				continue;
			}

			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<u64> {
		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<u8>) -> 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)
			},
		)
	})
}