From 1ffaf2669671ce18dede678ace266dbe6b4ff1bb Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 15 Oct 2026 15:42:50 +0000 Subject: [PATCH] Add a native federation store for the transmission loop `FederationStore` ports the store methods the federation transmission loop calls around each transaction to the native Postgres driver: `get_new_device_msgs_for_remote`, `delete_device_msgs_for_remote`, `get_device_updates_by_remote` (with the cross-signing key updates and cross-signatures), `mark_as_sent_devices_by_remote` and `set_destination_last_successful_stream_ordering`. The stream change cache checks those start with are left to the caller. None of the writes has a Python cache to invalidate. The cross-signing reads skip Python's caches. `tests/storage/test_federation_store.py` runs each against the Python method on one database, through a `FederationStoreForTest` pyclass. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01R5b7Ruy4cgt27Mjq6Ptxbb --- rust/src/federation/mod.rs | 1 + rust/src/storage/federation_store.rs | 917 +++++++++++++++++++++++++ rust/src/storage/mod.rs | 1 + synapse/synapse_rust/federation.pyi | 23 + tests/storage/test_federation_store.py | 367 ++++++++++ 5 files changed, 1309 insertions(+) create mode 100644 rust/src/storage/federation_store.rs create mode 100644 tests/storage/test_federation_store.py diff --git a/rust/src/federation/mod.rs b/rust/src/federation/mod.rs index 909d8c4471..10d6ab92ac 100644 --- a/rust/src/federation/mod.rs +++ b/rust/src/federation/mod.rs @@ -95,6 +95,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; m.add_submodule(&child_module)?; diff --git a/rust/src/storage/federation_store.rs b/rust/src/storage/federation_store.rs new file mode 100644 index 0000000000..1ac648cfbc --- /dev/null +++ b/rust/src/storage/federation_store.rs @@ -0,0 +1,917 @@ +/* + * This file is licensed under the Affero General Public License (AGPL) version 3. + * + * Copyright (C) 2026 Element Creations Ltd + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * See the GNU Affero General Public License for more details: + * . + * + */ + +//! What the federation transmission loop reads and writes, through the +//! native Postgres driver. +//! +//! Each method is a port of the store method it is named after, and +//! `tests/storage/test_federation_store.py` runs both against the same +//! database. The stream change cache checks those methods start with are +//! not here: the loop keeps its own, see `federation::queue::changes`. +//! +//! None of these writes has a Python cache to invalidate. + +use std::collections::HashMap; + +use anyhow::Context; +use base64::Engine; +use log::warn; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; +use serde_json::{Map, Value}; +use tokio::runtime::Handle; +use tokio_postgres::{types::ToSql, Row}; + +use crate::database::postgres::pool::{ConnectionPool, PooledConnection, PyConnectionPool}; +use crate::federation::queue::Edu; + +/// `EduTypes.DEVICE_LIST_UPDATE`. +const DEVICE_LIST_UPDATE: &str = "m.device_list_update"; +/// `EduTypes.SIGNING_KEY_UPDATE`. +const SIGNING_KEY_UPDATE: &str = "m.signing_key_update"; +/// The unstable name of `m.signing_key_update`, which is sent as well. +const UNSTABLE_SIGNING_KEY_UPDATE: &str = "org.matrix.signing_key_update"; + +/// The parts of the homeserver config the queries apply. +#[derive(Debug, Clone, Copy, Default)] +pub struct FederationStoreConfig { + /// `allow_device_name_lookup_over_federation`. + pub allow_device_name_lookup_over_federation: bool, +} + +/// The transmission loop's queries. +pub struct FederationStore { + config: FederationStoreConfig, +} + +impl FederationStore { + pub fn new(config: FederationStoreConfig) -> Self { + FederationStore { config } + } + + /// `get_new_device_msgs_for_remote`, less its stream change cache check: + /// the contents of the to-device messages queued for `destination` + /// after `last_stream_id`, and how far through the stream they go. + pub async fn get_new_device_msgs_for_remote( + &self, + pool: &ConnectionPool, + destination: &str, + last_stream_id: i64, + current_stream_id: i64, + limit: usize, + ) -> anyhow::Result<(Vec, i64)> { + if last_stream_id == current_stream_id { + return Ok((Vec::new(), current_stream_id)); + } + if limit == 0 { + // We ran out of room for EDUs in the transaction. + return Ok((Vec::new(), last_stream_id)); + } + + let conn = Conn::borrow(pool).await?; + let rows = conn + .query( + "SELECT stream_id, messages_json FROM device_federation_outbox + WHERE destination = $1 AND $2 < stream_id AND stream_id <= $3 + ORDER BY stream_id ASC + LIMIT $4", + &[ + &destination, + &last_stream_id, + ¤t_stream_id, + &(limit as i64), + ], + ) + .await?; + + let mut messages = Vec::with_capacity(rows.len()); + let mut stream_pos = current_stream_id; + for row in &rows { + stream_pos = row.try_get(0)?; + let json: &str = row.try_get(1)?; + messages.push(serde_json::from_str(json).context("parsing messages_json")?); + } + + // Short of the limit, there is nothing more up to the current + // position. + if messages.len() < limit { + stream_pos = current_stream_id; + } + Ok((messages, stream_pos)) + } + + /// `delete_device_msgs_for_remote`: the destination has the to-device + /// messages up to `up_to_stream_id`. + pub async fn delete_device_msgs_for_remote( + &self, + pool: &ConnectionPool, + destination: &str, + up_to_stream_id: i64, + ) -> anyhow::Result<()> { + Conn::borrow(pool) + .await? + .query( + "DELETE FROM device_federation_outbox + WHERE destination = $1 AND stream_id <= $2", + &[&destination, &up_to_stream_id], + ) + .await?; + Ok(()) + } + + /// `get_device_updates_by_remote`, less its stream change cache check: + /// the device list and cross-signing key updates to send to + /// `destination`, at most `limit`, and the stream ID they go up to. + /// + /// The EDUs of a user come in stream order, with the prev ID chain + /// through them. The users come in the order of their first update, + /// where Python has the order the database returns them in. + pub async fn get_device_updates_by_remote( + &self, + pool: &ConnectionPool, + destination: &str, + from_stream_id: i64, + now_stream_id: i64, + limit: usize, + ) -> anyhow::Result<(i64, Vec)> { + if from_stream_id == now_stream_id { + return Ok((now_stream_id, Vec::new())); + } + + let conn = Conn::borrow(pool).await?; + let rows = conn + .query( + "SELECT user_id, device_id, stream_id, opentracing_context + FROM device_lists_outbound_pokes + WHERE destination = $1 AND $2 < stream_id AND stream_id <= $3 + ORDER BY stream_id + LIMIT $4", + &[ + &destination, + &from_stream_id, + &now_stream_id, + &(limit as i64), + ], + ) + .await?; + if rows.is_empty() { + return Ok((now_stream_id, Vec::new())); + } + let updates = rows + .iter() + .map(|row| -> anyhow::Result { + Ok(Poke { + user_id: row.try_get(0)?, + device_id: row.try_get(1)?, + stream_id: row.try_get(2)?, + context: row.try_get(3)?, + }) + }) + .collect::>>()?; + + // The users' cross-signing keys, to tell which updates were of them. + let mut users: Vec<&str> = Vec::new(); + for update in &updates { + if !users.contains(&update.user_id.as_str()) { + users.push(&update.user_id); + } + } + let keys = cross_signing_keys(&conn, &users).await?; + let key_device = |user_id: &str, key_type: &str| -> anyhow::Result> { + keys.get(&(user_id.to_owned(), key_type.to_owned())) + .filter(|key| is_truthy(key)) + .map(cross_signing_key_device_id) + .transpose() + }; + let mut master_keys = HashMap::new(); + let mut self_signing_keys = HashMap::new(); + for user_id in &users { + if let Some(device_id) = key_device(user_id, "master")? { + master_keys.insert(*user_id, device_id); + } + if let Some(device_id) = key_device(user_id, "self_signing")? { + self_signing_keys.insert(*user_id, device_id); + } + } + + // The latest update per device, and the cross-signing key updates + // per user, within the budget. A cross-signing key update becomes + // two EDUs, so it costs two. + let mut last_processed_stream_id = from_stream_id; + let mut devices: OrderedDevices = OrderedDevices::default(); + let mut cross_signing: Vec<(&str, Map)> = Vec::new(); + for update in &updates { + let remaining = limit - (devices.len() + 2 * cross_signing.len()); + let user_id = update.user_id.as_str(); + let is_master = master_keys.get(user_id) == Some(&update.device_id); + let is_self_signing = self_signing_keys.get(user_id) == Some(&update.device_id); + let user_index = cross_signing.iter().position(|(user, _)| *user == user_id); + + if (is_master || is_self_signing) && user_index.is_none() && remaining < 2 { + break; + } + + if is_master || is_self_signing { + let index = user_index.unwrap_or_else(|| { + cross_signing.push((user_id, Map::new())); + cross_signing.len() - 1 + }); + let (field, key_type) = if is_master { + ("master_key", "master") + } else { + ("self_signing_key", "self_signing") + }; + let key = keys + .get(&(user_id.to_owned(), key_type.to_owned())) + .cloned() + .unwrap_or(Value::Null); + cross_signing[index].1.insert(field.to_owned(), key); + } else { + let key = (update.user_id.clone(), update.device_id.clone()); + if !devices.contains(&key) && remaining < 1 { + break; + } + devices.update(key, update.stream_id, update.context.clone()); + } + + last_processed_stream_id = update.stream_id; + } + + let mut edus = self + .device_update_edus(&conn, destination, from_stream_id, &devices) + .await?; + + for (user_id, mut content) in cross_signing { + content.insert("user_id".to_owned(), Value::from(user_id)); + let content = Value::Object(content); + edus.push(Edu { + edu_type: SIGNING_KEY_UPDATE.to_owned(), + content: content.clone(), + }); + edus.push(Edu { + edu_type: UNSTABLE_SIGNING_KEY_UPDATE.to_owned(), + content, + }); + } + + Ok((last_processed_stream_id, edus)) + } + + /// `_get_device_update_edus_by_remote`: an `m.device_list_update` EDU + /// per device, with its keys, or marked deleted. + async fn device_update_edus( + &self, + conn: &Conn, + destination: &str, + from_stream_id: i64, + devices: &OrderedDevices, + ) -> anyhow::Result> { + if devices.is_empty() { + return Ok(Vec::new()); + } + let found = device_keys(conn, devices.keys()).await?; + + let mut users: Vec<&str> = Vec::new(); + for (user_id, _) in devices.keys() { + if !users.contains(&user_id.as_str()) { + users.push(user_id); + } + } + let prev_ids = + last_device_updates_for_remote(conn, destination, &users, from_stream_id).await?; + + let mut edus = Vec::with_capacity(devices.len()); + for user_id in users { + // The prev ID of the first update is the last one sent before + // `from_stream_id`. + let mut prev_id = prev_ids.get(user_id).copied().unwrap_or(0); + + let mut user_devices: Vec<(&DeviceKey, &DeviceUpdate)> = devices + .entries + .iter() + .filter(|(key, _)| key.0 == user_id) + .map(|(key, value)| (key, value)) + .collect(); + user_devices.sort_by_key(|(_, (stream_id, _))| *stream_id); + + for (key, (stream_id, context)) in user_devices { + let mut content = Map::new(); + content.insert("user_id".to_owned(), Value::from(user_id)); + content.insert("device_id".to_owned(), Value::from(key.1.as_str())); + content.insert( + "prev_id".to_owned(), + if prev_id != 0 { + Value::from(vec![prev_id]) + } else { + Value::Array(Vec::new()) + }, + ); + content.insert("stream_id".to_owned(), Value::from(*stream_id)); + if context.as_deref() != Some("{}") { + content.insert( + "org.matrix.opentracing_context".to_owned(), + context.clone().map_or(Value::Null, Value::from), + ); + } + prev_id = *stream_id; + + match found.get(key) { + Some(device) => { + if let Some(keys) = device.keys.as_ref().filter(|keys| is_truthy(keys)) { + content.insert("keys".to_owned(), keys.clone()); + } + if self.config.allow_device_name_lookup_over_federation { + if let Some(name) = + device.display_name.as_ref().filter(|name| !name.is_empty()) + { + content.insert( + "device_display_name".to_owned(), + Value::from(name.as_str()), + ); + } + } + } + None => { + content.insert("deleted".to_owned(), Value::Bool(true)); + } + } + + edus.push(Edu { + edu_type: DEVICE_LIST_UPDATE.to_owned(), + content: Value::Object(content), + }); + } + } + Ok(edus) + } + + /// `mark_as_sent_devices_by_remote`: the destination has the device + /// list updates up to `stream_id`. Deletes the pokes and records, per + /// user, the last update sent, in one transaction. + pub async fn mark_as_sent_devices_by_remote( + &self, + pool: &ConnectionPool, + destination: &str, + stream_id: i64, + ) -> anyhow::Result<()> { + let mut conn = Conn::borrow(pool).await?; + conn.begin().await?; + + let rows = conn + .query( + "DELETE FROM device_lists_outbound_pokes + WHERE destination = $1 AND stream_id <= $2 + RETURNING user_id, stream_id", + &[&destination, &stream_id], + ) + .await?; + let mut max_stream_id_by_user: HashMap = HashMap::new(); + for row in &rows { + let user_id: String = row.try_get(0)?; + let poke_stream_id: i64 = row.try_get(1)?; + let max = max_stream_id_by_user.entry(user_id).or_insert(0); + *max = (*max).max(poke_stream_id); + } + + if !max_stream_id_by_user.is_empty() { + let (user_ids, stream_ids): (Vec, Vec) = + max_stream_id_by_user.into_iter().unzip(); + conn.query( + "INSERT INTO device_lists_outbound_last_success (destination, user_id, stream_id) + SELECT $1, user_id, stream_id FROM unnest($2::text[], $3::bigint[]) + AS u(user_id, stream_id) + ON CONFLICT (destination, user_id) DO UPDATE SET stream_id = EXCLUDED.stream_id", + &[&destination, &user_ids, &stream_ids], + ) + .await?; + } + + conn.commit().await + } + + /// `set_destination_last_successful_stream_ordering`. + pub async fn set_destination_last_successful_stream_ordering( + &self, + pool: &ConnectionPool, + destination: &str, + last_successful_stream_ordering: i64, + ) -> anyhow::Result<()> { + Conn::borrow(pool) + .await? + .query( + "INSERT INTO destinations (destination, last_successful_stream_ordering) + VALUES ($1, $2) + ON CONFLICT (destination) DO UPDATE + SET last_successful_stream_ordering = EXCLUDED.last_successful_stream_ordering", + &[&destination, &last_successful_stream_ordering], + ) + .await?; + Ok(()) + } +} + +/// A row of `device_lists_outbound_pokes`. +struct Poke { + user_id: String, + device_id: String, + stream_id: i64, + context: Option, +} + +/// `(user ID, device ID)`. +type DeviceKey = (String, String); + +/// `(stream ID, tracing context)` of a device's update. +type DeviceUpdate = (i64, Option); + +/// The latest update of each device, in the order the devices were first +/// seen. +#[derive(Default)] +struct OrderedDevices { + entries: Vec<(DeviceKey, DeviceUpdate)>, +} + +impl OrderedDevices { + fn len(&self) -> usize { + self.entries.len() + } + + fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + fn contains(&self, key: &(String, String)) -> bool { + self.entries.iter().any(|(k, _)| k == key) + } + + fn keys(&self) -> impl Iterator { + self.entries.iter().map(|(key, _)| key) + } + + /// Keep the update with the higher stream ID. + fn update(&mut self, key: (String, String), stream_id: i64, context: Option) { + match self.entries.iter_mut().find(|(k, _)| *k == key) { + Some((_, existing)) => { + if stream_id > existing.0 { + *existing = (stream_id, context); + } + } + None => self.entries.push((key, (stream_id, context))), + } + } +} + +/// A device's keys and name: `DeviceKeyLookupResult`. +struct DeviceKeys { + display_name: Option, + keys: Option, +} + +/// Whether Python would take `value` as true. +fn is_truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(b) => *b, + Value::Number(n) => n.as_f64() != Some(0.0), + Value::String(s) => !s.is_empty(), + Value::Array(a) => !a.is_empty(), + Value::Object(o) => !o.is_empty(), + } +} + +/// The latest cross-signing keys of each user, keyed by `(user ID, key +/// type)`: `_get_bare_e2e_cross_signing_keys_bulk`. +async fn cross_signing_keys( + conn: &Conn, + users: &[&str], +) -> anyhow::Result> { + let rows = conn + .query( + "SELECT DISTINCT ON (user_id, keytype) user_id, keytype, keydata + FROM e2e_cross_signing_keys + WHERE user_id = ANY($1) + ORDER BY user_id, keytype, stream_id DESC", + &[&users], + ) + .await?; + let mut keys = HashMap::with_capacity(rows.len()); + for row in &rows { + let user_id: String = row.try_get(0)?; + let key_type: String = row.try_get(1)?; + let key_data: &str = row.try_get(2)?; + keys.insert( + (user_id, key_type), + serde_json::from_str(key_data).context("parsing keydata")?, + ); + } + Ok(keys) +} + +/// The device ID of a cross-signing key: the part of its one key ID after +/// the algorithm. `get_verify_key_from_cross_signing_key` fails on anything +/// but a single valid ed25519 key, and so does this. +fn cross_signing_key_device_id(key: &Value) -> anyhow::Result { + let keys = key + .get("keys") + .and_then(Value::as_object) + .context("Invalid key")?; + let mut entries = keys.iter(); + let (Some((key_id, key_data)), None) = (entries.next(), entries.next()) else { + anyhow::bail!("Invalid key"); + }; + let (algorithm, version) = key_id.split_once(':').context("Invalid key")?; + if algorithm != "ed25519" { + anyhow::bail!("Unsupported algorithm {algorithm}"); + } + let key_data = key_data.as_str().context("Invalid key")?; + let bytes = base64::engine::general_purpose::STANDARD_NO_PAD + .decode(key_data.trim_end_matches('=')) + .context("Invalid key")?; + if bytes.len() != 32 { + anyhow::bail!("Invalid key"); + } + Ok(version.to_owned()) +} + +/// The devices' keys, with the owners' cross-signatures on them, and names: +/// `get_e2e_device_keys_and_signatures` with `include_all_devices`. Hidden +/// devices are left out. A device that is not in the result has been +/// deleted. +async fn device_keys( + conn: &Conn, + devices: impl Iterator, +) -> anyhow::Result> { + let (user_ids, device_ids): (Vec<&str>, Vec<&str>) = devices + .map(|(user_id, device_id)| (user_id.as_str(), device_id.as_str())) + .unzip(); + let rows = conn + .query( + "SELECT user_id, device_id, d.display_name, k.key_json + FROM devices d + LEFT JOIN e2e_device_keys_json k USING (user_id, device_id) + WHERE (user_id, device_id) IN ( + SELECT * FROM unnest($1::text[], $2::text[]) + ) AND NOT d.hidden", + &[&user_ids, &device_ids], + ) + .await?; + + let mut found = HashMap::with_capacity(rows.len()); + for row in &rows { + let user_id: String = row.try_get(0)?; + let device_id: String = row.try_get(1)?; + let display_name: Option = row.try_get(2)?; + let key_json: Option<&str> = row.try_get(3)?; + let keys = key_json + .filter(|json| !json.is_empty()) + .map(serde_json::from_str::) + .transpose() + .context("parsing key_json")?; + found.insert((user_id, device_id), DeviceKeys { display_name, keys }); + } + + // Add each owner's cross-signatures on the devices with keys. + let (user_ids, device_ids): (Vec<&str>, Vec<&str>) = found + .iter() + .filter(|(_, device)| device.keys.is_some()) + .map(|((user_id, device_id), _)| (user_id.as_str(), device_id.as_str())) + .unzip(); + if user_ids.is_empty() { + return Ok(found); + } + let rows = conn + .query( + "SELECT user_id, key_id, target_device_id, signature + FROM e2e_cross_signing_signatures + WHERE (target_user_id, target_device_id, user_id) IN ( + SELECT u, d, u FROM unnest($1::text[], $2::text[]) AS t(u, d) + )", + &[&user_ids, &device_ids], + ) + .await?; + let mut signatures: Vec<(String, String, String, String)> = Vec::with_capacity(rows.len()); + for row in &rows { + signatures.push(( + row.try_get(0)?, + row.try_get(1)?, + row.try_get(2)?, + row.try_get(3)?, + )); + } + for (user_id, key_id, device_id, signature) in signatures { + let Some(keys) = found + .get_mut(&(user_id.clone(), device_id)) + .and_then(|device| device.keys.as_mut()) + .and_then(Value::as_object_mut) + else { + continue; + }; + let by_user = keys + .entry("signatures") + .or_insert_with(|| Value::Object(Map::new())); + let Some(by_user) = by_user.as_object_mut() else { + continue; + }; + let by_key = by_user + .entry(user_id) + .or_insert_with(|| Value::Object(Map::new())); + if let Some(by_key) = by_key.as_object_mut() { + by_key.insert(key_id, Value::from(signature)); + } + } + Ok(found) +} + +/// The last device list update sent to `destination` for each user, at or +/// before `from_stream_id`: `_get_last_device_update_for_remote_user`, for +/// several users at once. A user with none is missing. +async fn last_device_updates_for_remote( + conn: &Conn, + destination: &str, + users: &[&str], + from_stream_id: i64, +) -> anyhow::Result> { + let rows = conn + .query( + "SELECT user_id, max(stream_id) FROM device_lists_outbound_last_success + WHERE destination = $1 AND user_id = ANY($2) AND stream_id <= $3 + GROUP BY user_id", + &[&destination, &users, &from_stream_id], + ) + .await?; + let mut last = HashMap::with_capacity(rows.len()); + for row in &rows { + let stream_id: Option = row.try_get(1)?; + last.insert(row.try_get(0)?, stream_id.unwrap_or(0)); + } + Ok(last) +} + +/// A connection borrowed for a group of statements, rolled back if it is +/// dropped inside a transaction. +struct Conn { + conn: Option, + in_transaction: bool, +} + +impl Conn { + async fn borrow(pool: &ConnectionPool) -> anyhow::Result { + let conn = pool + .get() + .await + .map_err(|e| anyhow::anyhow!("checking out a database connection: {e}"))?; + Ok(Conn { + conn: Some(conn), + in_transaction: false, + }) + } + + fn conn(&self) -> &PooledConnection { + self.conn.as_ref().expect("held until dropped") + } + + /// Run a statement through the connection's prepared statement cache. + async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> anyhow::Result> { + let conn = self.conn(); + let statement = conn + .prepare_cached(sql, false) + .await + .with_context(|| format!("preparing {}", sql.trim()))?; + conn.query(&statement, params) + .await + .with_context(|| format!("running {}", sql.trim())) + } + + /// Open a `READ COMMITTED` transaction. Its writes are upserts and + /// deletes on rows only this sender writes, so it needs no retry loop. + async fn begin(&mut self) -> anyhow::Result<()> { + self.conn() + .batch_execute("BEGIN ISOLATION LEVEL READ COMMITTED") + .await + .context("opening a transaction")?; + self.in_transaction = true; + Ok(()) + } + + async fn commit(mut self) -> anyhow::Result<()> { + self.conn() + .batch_execute("COMMIT") + .await + .context("committing")?; + self.in_transaction = false; + Ok(()) + } +} + +impl Drop for Conn { + fn drop(&mut self) { + let Some(conn) = self.conn.take() else { return }; + if !self.in_transaction { + return; + } + // Roll back before the connection goes back to the pool, so the next + // user doesn't find it inside our transaction. Without a runtime to + // do that on, the connection is dropped rather than returned. + match Handle::try_current() { + Ok(handle) => { + handle.spawn(async move { + if let Err(e) = conn.batch_execute("ROLLBACK").await { + warn!("Failed to roll back a federation transaction: {e}"); + } + drop(conn); + }); + } + Err(_) => { + drop(PooledConnection::take(conn)); + } + } + } +} + +/// The store, for tests that compare it with Python's on one database. +/// Each method blocks until its query is done. +#[pyclass(frozen, name = "FederationStoreForTest")] +pub struct PyFederationStoreForTest { + store: FederationStore, + pool: ConnectionPool, + handle: Handle, +} + +impl PyFederationStoreForTest { + fn run( + &self, + py: Python<'_>, + f: impl std::future::Future> + Send, + ) -> PyResult { + py.detach(|| self.handle.block_on(f)) + .map_err(|err| PyRuntimeError::new_err(format!("{err:#}"))) + } +} + +#[pymethods] +impl PyFederationStoreForTest { + #[new] + fn new( + pool: &Bound<'_, PyConnectionPool>, + allow_device_name_lookup_over_federation: bool, + ) -> Self { + let pool = pool.get(); + PyFederationStoreForTest { + store: FederationStore::new(FederationStoreConfig { + allow_device_name_lookup_over_federation, + }), + pool: pool.pool.clone(), + handle: pool.handle.clone(), + } + } + + fn get_new_device_msgs_for_remote<'py>( + &self, + py: Python<'py>, + destination: &str, + last_stream_id: i64, + current_stream_id: i64, + limit: usize, + ) -> PyResult<(Bound<'py, PyAny>, i64)> { + let (messages, stream_id) = self.run( + py, + self.store.get_new_device_msgs_for_remote( + &self.pool, + destination, + last_stream_id, + current_stream_id, + limit, + ), + )?; + Ok((pythonize::pythonize(py, &messages)?, stream_id)) + } + + fn delete_device_msgs_for_remote( + &self, + py: Python<'_>, + destination: &str, + up_to_stream_id: i64, + ) -> PyResult<()> { + self.run( + py, + self.store + .delete_device_msgs_for_remote(&self.pool, destination, up_to_stream_id), + ) + } + + #[allow(clippy::type_complexity)] + fn get_device_updates_by_remote<'py>( + &self, + py: Python<'py>, + destination: &str, + from_stream_id: i64, + now_stream_id: i64, + limit: usize, + ) -> PyResult<(i64, Vec<(String, Bound<'py, PyAny>)>)> { + let (stream_id, edus) = self.run( + py, + self.store.get_device_updates_by_remote( + &self.pool, + destination, + from_stream_id, + now_stream_id, + limit, + ), + )?; + let edus = edus + .into_iter() + .map(|edu| Ok((edu.edu_type, pythonize::pythonize(py, &edu.content)?))) + .collect::>()?; + Ok((stream_id, edus)) + } + + fn mark_as_sent_devices_by_remote( + &self, + py: Python<'_>, + destination: &str, + stream_id: i64, + ) -> PyResult<()> { + self.run( + py, + self.store + .mark_as_sent_devices_by_remote(&self.pool, destination, stream_id), + ) + } + + fn set_destination_last_successful_stream_ordering( + &self, + py: Python<'_>, + destination: &str, + stream_ordering: i64, + ) -> PyResult<()> { + self.run( + py, + self.store.set_destination_last_successful_stream_ordering( + &self.pool, + destination, + stream_ordering, + ), + ) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn device_id_of_a_cross_signing_key() { + let key = json!({ + "user_id": "@a:test", + "usage": ["master"], + "keys": {"ed25519:AbCdEfGhIjKlMnOpQrStUvWxYz0123456789+/AbCdE": "AbCdEfGhIjKlMnOpQrStUvWxYz0123456789+/AbCdE"}, + }); + assert_eq!( + cross_signing_key_device_id(&key).unwrap(), + "AbCdEfGhIjKlMnOpQrStUvWxYz0123456789+/AbCdE" + ); + + for bad in [ + json!({}), + json!({"keys": {}}), + json!({"keys": {"ed25519:a": "AAAA", "ed25519:b": "AAAA"}}), + json!({"keys": {"curve25519:a": "AbCdEfGhIjKlMnOpQrStUvWxYz0123456789+/AbCdE"}}), + json!({"keys": {"ed25519:a": "too short"}}), + ] { + assert!(cross_signing_key_device_id(&bad).is_err(), "{bad}"); + } + } + + #[test] + fn truthiness() { + assert!(!is_truthy(&json!({}))); + assert!(!is_truthy(&Value::Null)); + assert!(is_truthy(&json!({"a": 1}))); + } + + #[test] + fn latest_update_per_device() { + let mut devices = OrderedDevices::default(); + let a = ("@a:test".to_owned(), "A".to_owned()); + let b = ("@a:test".to_owned(), "B".to_owned()); + devices.update(a.clone(), 1, None); + devices.update(b.clone(), 2, Some("x".to_owned())); + devices.update(a.clone(), 3, Some("y".to_owned())); + assert_eq!(devices.len(), 2); + assert_eq!(devices.entries[0], (a, (3, Some("y".to_owned())))); + assert_eq!(devices.entries[1], (b, (2, Some("x".to_owned())))); + } +} diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index bbd4ae3a45..7756685db8 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -16,6 +16,7 @@ pub mod caches; pub mod db; pub mod events; +pub mod federation_store; pub mod python; pub mod store; pub mod sync_store; diff --git a/synapse/synapse_rust/federation.pyi b/synapse/synapse_rust/federation.pyi index 5d9019da16..59a238dafa 100644 --- a/synapse/synapse_rust/federation.pyi +++ b/synapse/synapse_rust/federation.pyi @@ -252,3 +252,26 @@ class RustDestinationQueue: def start_catching_up(self) -> None: ... def pending_pdu_ids(self) -> list[str]: """The event IDs of the queued PDUs, oldest first. For tests.""" + +class FederationStoreForTest: + """The Rust federation store, for tests that compare it with Python's. + Each method blocks until its query is done.""" + + def __init__( + self, pool: Any, allow_device_name_lookup_over_federation: bool + ) -> None: ... + def get_new_device_msgs_for_remote( + self, destination: str, last_stream_id: int, current_stream_id: int, limit: int + ) -> tuple[list[JsonDict], int]: ... + def delete_device_msgs_for_remote( + self, destination: str, up_to_stream_id: int + ) -> None: ... + def get_device_updates_by_remote( + self, destination: str, from_stream_id: int, now_stream_id: int, limit: int + ) -> tuple[int, list[tuple[str, JsonDict]]]: ... + def mark_as_sent_devices_by_remote( + self, destination: str, stream_id: int + ) -> None: ... + def set_destination_last_successful_stream_ordering( + self, destination: str, stream_ordering: int + ) -> None: ... diff --git a/tests/storage/test_federation_store.py b/tests/storage/test_federation_store.py new file mode 100644 index 0000000000..ecae32322a --- /dev/null +++ b/tests/storage/test_federation_store.py @@ -0,0 +1,367 @@ +# +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . +# + +"""The Rust federation store (`rust/src/storage/federation_store.rs`) against +the Python store methods it ports, on the same database. These need the native +Rust Postgres driver: + + SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_federation_store +""" + +from typing import Any +from unittest import SkipTest + +from signedjson.key import ( + encode_verify_key_base64, + generate_signing_key, + get_verify_key, +) + +from twisted.internet.testing import MemoryReactor + +from synapse.handlers.e2e_keys import SignatureListItem +from synapse.server import HomeServer +from synapse.storage.rust_pool import RustConnectionPool +from synapse.synapse_rust.federation import FederationStoreForTest +from synapse.types import JsonDict +from synapse.util.clock import Clock + +from tests import unittest + +DEST = "remote.example.com" +OTHER = "other.example.com" +ROOM = "!room:test" + + +def _sorted_edus(edus: list[tuple[str, JsonDict]]) -> list[tuple[str, JsonDict]]: + """The EDUs in a fixed order. Python orders the users as the database + returns them, and Rust by their first update.""" + return sorted( + edus, + key=lambda edu: (edu[0], edu[1]["user_id"], edu[1].get("stream_id", 0)), + ) + + +class FederationStoreTestCase(unittest.HomeserverTestCase): + """Each test runs a Rust query and the Python store method against the + same fixture and compares.""" + + def default_config(self) -> dict[str, Any]: + config = super().default_config() + # `get_device_updates_by_remote` only works on a federation sender. + config["federation_sender_instances"] = ["master"] + return config + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + hs = self.setup_test_homeserver() + if not isinstance( + hs.get_datastores().main.db_pool._db_pool, RustConnectionPool + ): + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastores().main + db_pool = self.store.db_pool._db_pool + assert isinstance(db_pool, RustConnectionPool) + assert db_pool._pool is not None + self.rust = FederationStoreForTest(db_pool._pool, False) + self.rust_with_names = FederationStoreForTest(db_pool._pool, True) + + # To-device messages. + + def _to_device_edu(self, n: int) -> JsonDict: + return { + "sender": "@alice:test", + "type": "m.test", + "message_id": str(n), + "messages": {"@u:remote": {"D": {"n": n}}}, + } + + def _queue_to_device(self, destination: str, n: int) -> None: + self.get_success( + self.store.add_remote_messages_from_client_to_device_inbox( + {destination: self._to_device_edu(n)} + ) + ) + + def test_new_device_msgs_for_remote(self) -> None: + for n in range(3): + self._queue_to_device(DEST, n) + self._queue_to_device(OTHER, n) + current = self.store.get_to_device_stream_token() + + for last in (0, current - 3, current): + for limit in (0, 1, 2, 10): + self.assertEqual( + self.rust.get_new_device_msgs_for_remote( + DEST, last, current, limit + ), + self.get_success( + self.store.get_new_device_msgs_for_remote( + DEST, last, current, limit + ) + ), + (last, limit), + ) + + def test_delete_device_msgs_for_remote(self) -> None: + for n in range(3): + self._queue_to_device(DEST, n) + self._queue_to_device(OTHER, 0) + current = self.store.get_to_device_stream_token() + + self.rust.delete_device_msgs_for_remote(DEST, current - 2) + + messages, _ = self.get_success( + self.store.get_new_device_msgs_for_remote(DEST, 0, current, 10) + ) + self.assertEqual(messages, [self._to_device_edu(2)]) + messages, _ = self.get_success( + self.store.get_new_device_msgs_for_remote(OTHER, 0, current, 10) + ) + self.assertEqual(len(messages), 1) + + # Device list updates. + + def _poke(self, user_id: str, device_id: str, host: str = DEST) -> None: + self.get_success( + self.store.add_device_change_to_streams(user_id, [device_id], [ROOM]) + ) + self.get_success( + self.store.add_device_list_outbound_pokes( + user_id=user_id, + device_id=device_id, + room_id=ROOM, + hosts=[host], + context={}, + ) + ) + + def _cross_signing_key(self, user_id: str, key_type: str) -> str: + """Give the user a cross-signing key, and return its device ID.""" + public = encode_verify_key_base64(get_verify_key(generate_signing_key("x"))) + key = { + "user_id": user_id, + "usage": [key_type], + "keys": {f"ed25519:{public}": public}, + } + self.get_success(self.store.set_e2e_cross_signing_key(user_id, key_type, key)) + return public + + def _device_keys(self, user_id: str, device_id: str) -> None: + self.get_success( + self.store.set_e2e_device_keys( + user_id, + device_id, + self.clock.time_msec(), + { + "user_id": user_id, + "device_id": device_id, + "algorithms": ["m.olm.v1.curve25519-aes-sha2"], + "keys": {f"ed25519:{device_id}": "key"}, + "signatures": {user_id: {f"ed25519:{device_id}": "self"}}, + }, + ) + ) + + def _device_list_fixture(self) -> None: + """Pokes to DEST for: devices with and without keys and names, a + device's repeated update, a deleted device, a hidden device, a + cross-signed device, and both kinds of cross-signing key update. And + one poke for another destination.""" + alice = "@alice:test" + bob = "@bob:test" + + self.get_success(self.store.store_device(alice, "A1", "Alice's phone")) + self.get_success(self.store.store_device(alice, "A2", None)) + self.get_success(self.store.store_device(bob, "B1", "Bob's laptop")) + self._device_keys(alice, "A1") + self._device_keys(bob, "B1") + + master = self._cross_signing_key(alice, "master") + self_signing = self._cross_signing_key(alice, "self_signing") + self.get_success( + self.store.store_e2e_cross_signing_signatures( + alice, + [ + SignatureListItem( + signing_key_id=f"ed25519:{self_signing}", + target_user_id=alice, + target_device_id="A1", + signature="cross", # type: ignore[arg-type] + ) + ], + ) + ) + + self._poke(alice, "A1") + self._poke(bob, "B1") + self._poke(alice, "A2") + self._poke(alice, master) + self._poke(alice, "A1") + self._poke(bob, "GONE") + self._poke(alice, self_signing) + self._poke(bob, "B1", host=OTHER) + + # A hidden device is sent as deleted. + self.get_success(self.store.store_device(bob, "HIDDEN", None)) + self.get_success( + self.store.db_pool.simple_update_one( + "devices", {"user_id": bob, "device_id": "HIDDEN"}, {"hidden": True} + ) + ) + self._poke(bob, "HIDDEN") + + def _assert_device_updates(self, from_stream_id: int, limit: int) -> None: + now = self.store.get_device_stream_token().stream + python = self.get_success( + self.store.get_device_updates_by_remote(DEST, from_stream_id, limit) + ) + for rust, names in ((self.rust, False), (self.rust_with_names, True)): + self.hs.config.federation.allow_device_name_lookup_over_federation = names + python = self.get_success( + self.store.get_device_updates_by_remote(DEST, from_stream_id, limit) + ) + stream_id, edus = rust.get_device_updates_by_remote( + DEST, from_stream_id, now, limit + ) + self.assertEqual(stream_id, python[0], (from_stream_id, limit, names)) + self.assertEqual( + _sorted_edus(edus), + _sorted_edus(python[1]), + (from_stream_id, limit, names), + ) + + def test_device_updates_by_remote(self) -> None: + self._device_list_fixture() + for limit in (1, 2, 3, 100): + self._assert_device_updates(0, limit) + + # The fixture has every kind of update. + _, edus = self.rust_with_names.get_device_updates_by_remote( + DEST, 0, self.store.get_device_stream_token().stream, 100 + ) + types = [edu_type for edu_type, _ in edus] + self.assertEqual(types.count("m.signing_key_update"), 1) + self.assertEqual(types.count("org.matrix.signing_key_update"), 1) + updates = { + content["device_id"]: content + for edu_type, content in edus + if edu_type == "m.device_list_update" + } + self.assertEqual(set(updates), {"A1", "A2", "B1", "GONE", "HIDDEN"}) + self.assertEqual(updates["A1"]["device_display_name"], "Alice's phone") + self.assertIn("cross", str(updates["A1"]["keys"]["signatures"])) + self.assertNotIn("keys", updates["A2"]) + self.assertTrue(updates["GONE"]["deleted"]) + self.assertTrue(updates["HIDDEN"]["deleted"]) + # A1 was updated again after A2, so it follows A2. + self.assertEqual(updates["A1"]["prev_id"], [updates["A2"]["stream_id"]]) + + def test_device_updates_by_remote_paginates(self) -> None: + """Walk the stream a few updates at a time, marking each batch sent, + so the prev IDs come from the last batch.""" + self._device_list_fixture() + position = 0 + for _ in range(10): + self._assert_device_updates(position, 3) + stream_id, edus = self.rust.get_device_updates_by_remote( + DEST, position, self.store.get_device_stream_token().stream, 3 + ) + if not edus: + break + self.rust.mark_as_sent_devices_by_remote(DEST, stream_id) + position = stream_id + else: + self.fail("never caught up") + + def test_mark_as_sent_devices_by_remote(self) -> None: + """The pokes sent are deleted, and each user's last one recorded.""" + self._device_list_fixture() + stream_id, _ = self.get_success( + self.store.get_device_updates_by_remote(DEST, 0, limit=4) + ) + pokes = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_pokes", + {"destination": DEST}, + ["user_id", "stream_id"], + ) + ) + expected: dict[str, int] = {} + for user_id, poke_stream_id in pokes: + if poke_stream_id <= stream_id: + expected[user_id] = max(expected.get(user_id, 0), poke_stream_id) + + self.rust.mark_as_sent_devices_by_remote(DEST, stream_id) + + remaining = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_pokes", + {"destination": DEST}, + ["stream_id"], + ) + ) + self.assertTrue(remaining) + self.assertTrue(all(row[0] > stream_id for row in remaining)) + last_success = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_last_success", + {"destination": DEST}, + ["user_id", "stream_id"], + ) + ) + self.assertEqual(dict(last_success), expected) + # Another destination's pokes are left alone. + other = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_pokes", {"destination": OTHER}, ["stream_id"] + ) + ) + self.assertEqual(len(other), 1) + + # Marking again moves a user's last success on. + self.rust.mark_as_sent_devices_by_remote( + DEST, self.store.get_device_stream_token().stream + ) + last_success = self.get_success( + self.store.db_pool.simple_select_list( + "device_lists_outbound_last_success", + {"destination": DEST}, + ["user_id", "stream_id"], + ) + ) + self.assertEqual( + {user_id for user_id, _ in last_success}, {"@alice:test", "@bob:test"} + ) + self.assertTrue( + all(stream > expected.get(user_id, 0) for user_id, stream in last_success) + ) + + def test_set_destination_last_successful_stream_ordering(self) -> None: + self.rust.set_destination_last_successful_stream_ordering(DEST, 5) + self.assertEqual( + self.get_success( + self.store.get_destination_last_successful_stream_ordering(DEST) + ), + 5, + ) + self.rust.set_destination_last_successful_stream_ordering(DEST, 7) + self.assertEqual( + self.get_success( + self.store.get_destination_last_successful_stream_ordering(DEST) + ), + 7, + )