tuwunel_core/matrix/pdu/

unsigned.rs

use std::collections::BTreeMap;

use ruma::{
	MilliSecondsSinceUnixEpoch, OwnedEventId,
	events::{AnySyncMessageLikeEvent, room::member::MembershipState},
	serde::Raw,
};
use serde::Serialize;
use serde_json::value::{RawValue as RawJsonValue, Value as JsonValue, to_raw_value};

use super::{Pdu, Unsigned};
use crate::{Result, err, implement};

#[implement(Pdu)]
pub fn remove_transaction_id(&mut self) -> Result {
	use BTreeMap as Map;

	let Some(unsigned) = &self.unsigned else {
		return Ok(());
	};

	let mut unsigned: Map<&str, Raw<JsonValue>> = serde_json::from_str(unsigned.json().get())
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	unsigned.remove("transaction_id");
	self.unsigned = to_raw_value(&unsigned)
		.map(Into::into)
		.map(Some)
		.expect("unsigned is valid");

	Ok(())
}

#[implement(Pdu)]
pub fn add_age(&mut self) -> Result {
	use BTreeMap as Map;

	let mut unsigned: Map<&str, Raw<JsonValue>> = self
		.unsigned
		.as_ref()
		.map(Unsigned::json)
		.map(RawJsonValue::get)
		.map_or_else(|| Ok(Map::new()), serde_json::from_str)
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	// deliberately allowing for the possibility of negative age
	let now: i128 = MilliSecondsSinceUnixEpoch::now().get().into();
	let then: i128 = self.origin_server_ts.into();
	let this_age = now.saturating_sub(then);

	unsigned.insert("age", raw_of(&this_age)?);
	self.unsigned = Some(to_raw_value(&unsigned)?.into());

	Ok(())
}

/// MSC4115: annotate the served event with the requesting user's room
/// membership at the time of the event.
#[implement(Pdu)]
pub fn add_membership(&mut self, membership: &MembershipState) -> Result {
	use BTreeMap as Map;

	let mut unsigned: Map<&str, Raw<JsonValue>> = self
		.unsigned
		.as_ref()
		.map(Unsigned::json)
		.map(RawJsonValue::get)
		.map_or_else(|| Ok(Map::new()), serde_json::from_str)
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	unsigned.insert("membership", raw_of(membership)?);
	self.unsigned = Some(to_raw_value(&unsigned)?.into());

	Ok(())
}

#[implement(Pdu)]
pub fn add_relation(&mut self, name: &str, pdu: Option<&Pdu>) -> Result {
	use serde_json::Map;

	let mut unsigned: Map<String, JsonValue> = self
		.unsigned
		.as_ref()
		.map(Unsigned::json)
		.map(RawJsonValue::get)
		.map_or_else(|| Ok(Map::new()), serde_json::from_str)
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	let pdu = pdu
		.map(serde_json::to_value)
		.transpose()?
		.unwrap_or_else(|| JsonValue::Object(Map::new()));

	unsigned
		.entry("m.relations")
		.or_insert(JsonValue::Object(Map::new()))
		.as_object_mut()
		.map(|object| object.insert(name.to_owned(), pdu));

	self.unsigned = Some(to_raw_value(&unsigned)?.into());

	Ok(())
}

/// MSC3816: overwrite `unsigned.m.relations.m.thread.current_user_participated`
/// with a per-requester value. No-op when the event carries no thread bundle.
#[implement(Pdu)]
pub fn set_thread_participated(&mut self, participated: bool) -> Result {
	use serde_json::Map;

	let Some(unsigned) = self.unsigned.as_ref() else {
		return Ok(());
	};

	let mut unsigned: Map<String, JsonValue> = serde_json::from_str(unsigned.json().get())
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	let updated = unsigned
		.get_mut("m.relations")
		.and_then(JsonValue::as_object_mut)
		.and_then(|relations| relations.get_mut("m.thread"))
		.and_then(JsonValue::as_object_mut)
		.map(|thread| {
			thread.insert("current_user_participated".to_owned(), participated.into());
		})
		.is_some();

	if updated {
		self.unsigned = Some(to_raw_value(&unsigned)?.into());
	}

	Ok(())
}

/// MSC3925: fold the newest `m.replace` edit into
/// `unsigned.m.relations.m.replace` as the full replacement event, preserving
/// an existing bundle such as `m.thread` and creating `unsigned` when absent.
#[implement(Pdu)]
pub fn set_replacement_bundle(&mut self, replacement: &Raw<AnySyncMessageLikeEvent>) -> Result {
	use BTreeMap as Map;

	type Object = Map<String, Raw<JsonValue>>;

	let parse = |raw: &RawJsonValue| -> Result<Object> {
		serde_json::from_str(raw.get())
			.map_err(|e| err!(Database("Invalid object in pdu unsigned: {e}")))
	};

	let mut unsigned: Object = self
		.unsigned
		.as_ref()
		.map(|unsigned| parse(unsigned.json()))
		.transpose()?
		.unwrap_or_default();

	let mut relations: Object = unsigned
		.get("m.relations")
		.map(|relations| parse(relations.json()))
		.transpose()?
		.unwrap_or_default();

	relations.insert("m.replace".to_owned(), replacement.cast_ref().clone());
	unsigned.insert("m.relations".to_owned(), to_raw_value(&relations)?.into());
	self.unsigned = Some(to_raw_value(&unsigned)?.into());

	Ok(())
}

/// MSC2675/MSC3267: fold reference relations into
/// `unsigned.m.relations.m.reference` as `{ chunk: [{ event_id }, ...] }`,
/// preserving an existing bundle such as `m.thread` or `m.replace` and creating
/// `unsigned` when absent.
#[implement(Pdu)]
pub fn set_reference_bundle(&mut self, event_ids: &[OwnedEventId]) -> Result {
	use BTreeMap as Map;

	type Object = Map<String, Raw<JsonValue>>;

	let parse = |raw: &RawJsonValue| -> Result<Object> {
		serde_json::from_str(raw.get())
			.map_err(|e| err!(Database("Invalid object in pdu unsigned: {e}")))
	};

	let mut unsigned: Object = self
		.unsigned
		.as_ref()
		.map(|unsigned| parse(unsigned.json()))
		.transpose()?
		.unwrap_or_default();

	let mut relations: Object = unsigned
		.get("m.relations")
		.map(|relations| parse(relations.json()))
		.transpose()?
		.unwrap_or_default();

	let chunk: Vec<JsonValue> = event_ids
		.iter()
		.map(|event_id| serde_json::json!({ "event_id": event_id }))
		.collect();

	let reference = serde_json::json!({ "chunk": chunk });

	relations.insert("m.reference".to_owned(), to_raw_value(&reference)?.into());
	unsigned.insert("m.relations".to_owned(), to_raw_value(&relations)?.into());
	self.unsigned = Some(to_raw_value(&unsigned)?.into());

	Ok(())
}

#[inline]
fn raw_of<T: Serialize>(value: &T) -> Result<Raw<JsonValue>> {
	Ok(Raw::from_raw_value(&to_raw_value(value)?))
}