continuwuity/src/core/pdu/unsigned.rs

use std::collections::BTreeMap;

use ruma::MilliSecondsSinceUnixEpoch;
use serde::Deserialize;
use serde_json::value::{to_raw_value, RawValue as RawJsonValue, Value as JsonValue};

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

#[implement(Pdu)]
pub fn remove_transaction_id(&mut self) -> Result {
	let Some(unsigned) = &self.unsigned else {
		return Ok(());
	};

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

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

	Ok(())
}

#[implement(Pdu)]
pub fn add_age(&mut self) -> Result {
	let mut unsigned: BTreeMap<String, Box<RawJsonValue>> = self
		.unsigned
		.as_ref()
		.map_or_else(|| Ok(BTreeMap::new()), |u| serde_json::from_str(u.get()))
		.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".to_owned(), to_raw_value(&this_age).expect("age is valid"));
	self.unsigned = to_raw_value(&unsigned)
		.map(Some)
		.expect("unsigned is valid");

	Ok(())
}

#[implement(Pdu)]
pub fn add_relation(&mut self, name: &str, pdu: &Pdu) -> Result {
	let mut unsigned: BTreeMap<String, JsonValue> = self
		.unsigned
		.as_ref()
		.map_or_else(|| Ok(BTreeMap::new()), |u| serde_json::from_str(u.get()))
		.map_err(|e| err!(Database("Invalid unsigned in pdu event: {e}")))?;

	let relations: &mut JsonValue = unsigned.entry("m.relations".into()).or_default();
	if relations.as_object_mut().is_none() {
		let mut object = serde_json::Map::<String, JsonValue>::new();
		_ = relations.as_object_mut().insert(&mut object);
	}

	relations
		.as_object_mut()
		.expect("we just created it")
		.insert(name.to_owned(), serde_json::to_value(pdu)?);

	self.unsigned = to_raw_value(&unsigned)
		.map(Some)
		.expect("unsigned is valid");

	Ok(())
}

#[implement(Pdu)]
pub fn contains_unsigned_property<F>(&self, property: &str, is_type: F) -> bool
where
	F: FnOnce(&JsonValue) -> bool,
{
	self.get_unsigned_as_value()
		.get(property)
		.map(is_type)
		.is_some_and(is_true!())
}

#[implement(Pdu)]
pub fn get_unsigned_property<T>(&self, property: &str) -> Result<T>
where
	T: for<'de> Deserialize<'de>,
{
	self.get_unsigned_as_value()
		.get_mut(property)
		.map(JsonValue::take)
		.map(serde_json::from_value)
		.ok_or(err!(Request(NotFound("property not found in unsigned object"))))?
		.map_err(|e| err!(Database("Failed to deserialize unsigned.{property} into type: {e}")))
}

#[implement(Pdu)]
#[must_use]
pub fn get_unsigned_as_value(&self) -> JsonValue { self.get_unsigned::<JsonValue>().unwrap_or_default() }

#[implement(Pdu)]
pub fn get_unsigned<T>(&self) -> Result<JsonValue> {
	self.unsigned
		.as_ref()
		.map(|raw| raw.get())
		.map(serde_json::from_str)
		.ok_or(err!(Request(NotFound("\"unsigned\" property not found in pdu"))))?
		.map_err(|e| err!(Database("Failed to deserialize \"unsigned\" into value: {e}")))
}