fix: SEC8

Reviewed-By: Ginger <ginger@gingershaped.computer>
Reviewed-By: Jade <jade@ellis.link>
Reviewed-By: Jacob Taylor <jacob@explodie.org>

src/service/server_keys/acquire.rs

@@ -1,248 +0,0 @@
-use std::{
-	borrow::Borrow,
-	collections::{BTreeMap, BTreeSet},
-	time::Duration,
-};
-
-use conduwuit::{debug, debug_error, debug_warn, error, info, result::FlatOk, trace, warn};
-use futures::{StreamExt, stream::FuturesUnordered};
-use ruma::{
-	CanonicalJsonObject, OwnedServerName, OwnedServerSigningKeyId, ServerName,
-	ServerSigningKeyId, api::federation::discovery::ServerSigningKeys, serde::Raw,
-};
-use serde_json::value::RawValue as RawJsonValue;
-use tokio::time::{Instant, timeout_at};
-
-use super::key_exists;
-
-type Batch = BTreeMap<OwnedServerName, Vec<OwnedServerSigningKeyId>>;
-
-impl super::Service {
-	/// Returns all verify keys required to verify the incoming events.
-	pub async fn acquire_events_pubkeys<'a, I>(&self, events: I)
-	where
-		I: Iterator<Item = &'a Box<RawJsonValue>> + Send,
-	{
-		type Batch = BTreeMap<OwnedServerName, BTreeSet<OwnedServerSigningKeyId>>;
-		type Signatures = BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, String>>;
-
-		let mut batch = Batch::new();
-		events
-			.cloned()
-			.map(Raw::<CanonicalJsonObject>::from_json)
-			.map(|event| event.get_field::<Signatures>("signatures"))
-			.filter_map(FlatOk::flat_ok)
-			.flat_map(IntoIterator::into_iter)
-			.for_each(|(server, sigs)| {
-				batch.entry(server).or_default().extend(sigs.into_keys());
-			});
-
-		let batch = batch
-			.iter()
-			.map(|(server, keys)| (server.borrow(), keys.iter().map(Borrow::borrow)));
-
-		self.acquire_pubkeys(batch).await;
-	}
-
-	/// Acquires all the verify keys required in the provided batch, using
-	/// notaries and direct-fetch where applicable.
-	pub async fn acquire_pubkeys<'a, S, K>(&self, batch: S)
-	where
-		S: Iterator<Item = (&'a ServerName, K)> + Send + Clone,
-		K: Iterator<Item = &'a ServerSigningKeyId> + Send + Clone,
-	{
-		let notary_only = self.services.server.config.only_query_trusted_key_servers;
-		let notary_first_always = self.services.server.config.query_trusted_key_servers_first;
-		let notary_first_on_join = self
-			.services
-			.server
-			.config
-			.query_trusted_key_servers_first_on_join;
-
-		let requested_servers = batch.clone().count();
-		let requested_keys = batch.clone().flat_map(|(_, key_ids)| key_ids).count();
-
-		debug!("acquire {requested_keys} keys from {requested_servers}");
-
-		let mut missing = self.acquire_locals(batch).await;
-		let mut missing_keys = keys_count(&missing);
-		let mut missing_servers = missing.len();
-		if missing_servers == 0 {
-			return;
-		}
-
-		info!("{missing_keys} keys for {missing_servers} servers will be acquired");
-
-		if notary_first_always || notary_first_on_join {
-			missing = self.acquire_notary(missing.into_iter()).await;
-			missing_keys = keys_count(&missing);
-			missing_servers = missing.len();
-			if missing_keys == 0 {
-				return;
-			}
-
-			warn!(
-				"missing {missing_keys} keys for {missing_servers} servers from all notaries \
-				 first"
-			);
-		}
-
-		if !notary_only {
-			missing = self.acquire_origins(missing.into_iter()).await;
-			missing_keys = keys_count(&missing);
-			missing_servers = missing.len();
-			if missing_keys == 0 {
-				return;
-			}
-
-			debug_warn!("missing {missing_keys} keys for {missing_servers} servers unreachable");
-		}
-
-		if !notary_first_always && !notary_first_on_join {
-			missing = self.acquire_notary(missing.into_iter()).await;
-			missing_keys = keys_count(&missing);
-			missing_servers = missing.len();
-			if missing_keys == 0 {
-				return;
-			}
-
-			debug_warn!(
-				"still missing {missing_keys} keys for {missing_servers} servers from all \
-				 notaries."
-			);
-		}
-
-		if missing_keys > 0 {
-			warn!(
-				"did not obtain {missing_keys} keys for {missing_servers} servers out of \
-				 {requested_keys} total keys for {requested_servers} total servers."
-			);
-		}
-
-		for (server, key_ids) in missing {
-			debug_warn!(%server, ?key_ids, "missing");
-		}
-	}
-
-	/// Filters the batch to return only keys which we do not already have
-	/// locally.
-	async fn acquire_locals<'a, S, K>(&self, batch: S) -> Batch
-	where
-		S: Iterator<Item = (&'a ServerName, K)> + Send,
-		K: Iterator<Item = &'a ServerSigningKeyId> + Send,
-	{
-		let mut missing = Batch::new();
-		for (server, key_ids) in batch {
-			for key_id in key_ids {
-				if !self.verify_key_exists(server, key_id).await {
-					missing
-						.entry(server.into())
-						.or_default()
-						.push(key_id.into());
-				}
-			}
-		}
-
-		missing
-	}
-
-	/// Fetches verification keys from the origin servers directly, returning
-	/// any keys that could not be fetched.
-	async fn acquire_origins<I>(&self, batch: I) -> Batch
-	where
-		I: Iterator<Item = (OwnedServerName, Vec<OwnedServerSigningKeyId>)> + Send,
-	{
-		let timeout = Instant::now()
-			.checked_add(Duration::from_secs(45))
-			.expect("timeout overflows");
-
-		let mut requests: FuturesUnordered<_> = batch
-			.map(|(origin, key_ids)| self.acquire_origin(origin, key_ids, timeout))
-			.collect();
-
-		let mut missing = Batch::new();
-		while let Some((origin, key_ids)) = requests.next().await {
-			if !key_ids.is_empty() {
-				missing.insert(origin, key_ids);
-			}
-		}
-
-		missing
-	}
-
-	/// Fetches verification keys from an origin directly, returning any keys
-	/// that were fetched.
-	async fn acquire_origin(
-		&self,
-		origin: OwnedServerName,
-		mut key_ids: Vec<OwnedServerSigningKeyId>,
-		timeout: Instant,
-	) -> (OwnedServerName, Vec<OwnedServerSigningKeyId>) {
-		match timeout_at(timeout, self.server_request(&origin)).await {
-			| Err(e) => debug_warn!(%origin, "timed out: {e}"),
-			| Ok(Err(e)) => debug_error!(%origin, "{e}"),
-			| Ok(Ok(server_keys)) => {
-				trace!(
-					%origin,
-					?key_ids,
-					?server_keys,
-					"received server_keys"
-				);
-
-				self.add_signing_keys(server_keys.clone()).await;
-				key_ids.retain(|key_id| !key_exists(&server_keys, key_id));
-			},
-		}
-
-		(origin, key_ids)
-	}
-
-	/// Fetches signing keys for the given batch from the configured notary
-	/// servers, returning any keys that could not be fetched.
-	async fn acquire_notary<I>(&self, batch: I) -> Batch
-	where
-		I: Iterator<Item = (OwnedServerName, Vec<OwnedServerSigningKeyId>)> + Send,
-	{
-		let mut missing: Batch = batch.collect();
-		for notary in self.services.globals.trusted_servers() {
-			let missing_keys = keys_count(&missing);
-			let missing_servers = missing.len();
-			debug!(
-				"Asking notary {notary} for {missing_keys} missing keys from {missing_servers} \
-				 servers"
-			);
-
-			let batch = missing
-				.iter()
-				.map(|(server, keys)| (server.borrow(), keys.iter().map(Borrow::borrow)));
-
-			match self.batch_notary_request(notary, batch).await {
-				| Err(e) => error!("Failed to contact notary {notary:?}: {e}"),
-				| Ok(results) =>
-					for server_keys in results {
-						self.acquire_notary_result(&mut missing, server_keys).await;
-					},
-			}
-		}
-
-		missing
-	}
-
-	/// Contacts a notary and fetches signing keys for the given batch, removing
-	/// keys from the missing batch when they are found.
-	async fn acquire_notary_result(&self, missing: &mut Batch, server_keys: ServerSigningKeys) {
-		let server = &server_keys.server_name;
-		self.add_signing_keys(server_keys.clone()).await;
-
-		if let Some(key_ids) = missing.get_mut(server) {
-			key_ids.retain(|key_id| !key_exists(&server_keys, key_id));
-			if key_ids.is_empty() {
-				missing.remove(server);
-			}
-		}
-	}
-}
-
-fn keys_count(batch: &Batch) -> usize {
-	batch.values().flat_map(|key_ids| key_ids.iter()).count()
-}

src/service/server_keys/get.rs

@@ -1,146 +1,526 @@
-use std::borrow::Borrow;
+/// get.rs handles actually acquiring keys from remote servers.
+use std::collections::BTreeMap;
 
-use conduwuit::{Err, Result, trace};
+use assign::assign;
+use conduwuit::{Result, debug, debug_error, debug_info, debug_warn, err, info, trace};
+use futures::{StreamExt, stream::FuturesUnordered};
 use ruma::{
-	CanonicalJsonObject, ServerName, ServerSigningKeyId, api::federation::discovery::VerifyKey,
+	CanonicalJsonObject, CanonicalJsonValue, MilliSecondsSinceUnixEpoch, ServerName,
+	ServerSigningKeyId, UInt,
+	api::federation::discovery::{
+		ServerSigningKeys, VerifyKey, get_remote_server_keys_batch::v2::QueryCriteria,
+	},
 	room_version_rules::RoomVersionRules,
+	serde::Base64,
+	uint,
 };
+use serde_json::value::RawValue;
+use tokio::select;
 
-use super::{PubKeyMap, PubKeys, extract_key};
-use crate::server_keys::util::required_keys;
+use super::{PubKeyMap, in_one_week};
+use crate::server_keys::{request::ServerKeysQuery, util::required_keys};
 
 impl super::Service {
+	/// Extracts the minimum timestamp within which a signing key must be valid
+	/// for it to be able to verify this event.
+	pub(super) fn min_valid_ts_from_event(
+		object: &CanonicalJsonObject,
+	) -> Result<MilliSecondsSinceUnixEpoch> {
+		object
+			.get("origin_server_ts")
+			.and_then(CanonicalJsonValue::as_integer)
+			.map(i64::from)
+			.map(|ts| {
+				MilliSecondsSinceUnixEpoch(UInt::new_saturating(
+					u64::try_from(ts).unwrap_or_default(),
+				))
+			})
+			.ok_or_else(|| err!(Request(BadJson("Missing origin_server_ts in event"))))
+	}
+
 	/// Fetches the verify keys required to verify the incoming event.
+	///
+	/// Only returns an error if the `min_valid_ts` cannot be inferred from the
+	/// object, or if calculating which keys are required fails (malformed
+	/// event).
+	///
+	/// Keys which cannot be fetched are not included in the resulting map, but
+	/// will not cause an error.
 	pub async fn get_event_keys(
 		&self,
 		object: &CanonicalJsonObject,
 		room_version_rules: &RoomVersionRules,
 	) -> Result<PubKeyMap> {
-		let required = match required_keys(object, &room_version_rules.signatures) {
-			| Ok(required) => required,
-			| Err(e) => {
-				return Err!(BadServerResponse(
-					"Failed to determine keys required to verify event: {e}"
-				));
-			},
-		};
+		let min_valid_ts = Self::min_valid_ts_from_event(object)?;
+
+		let required = required_keys(object, &room_version_rules.signatures).map_err(|e| {
+			err!(BadServerResponse("Failed to determine keys required to verify event: {e}"))
+		})?;
 
 		let batch = required
-			.iter()
-			.map(|(s, ids)| (s.borrow(), ids.iter().map(Borrow::borrow)));
+			.into_iter()
+			.map(|(s, ids)| {
+				(
+					s,
+					ids.into_iter()
+						.map(|key_id| (key_id, new_query_criteria(min_valid_ts)))
+						.collect::<BTreeMap<_, _>>(),
+				)
+			})
+			.collect::<ServerKeysQuery>();
 
-		Ok(self.get_pubkeys(batch).await)
+		Ok(self.fetch_verify_keys(batch).await)
 	}
 
-	/// Fetches all verify keys for the given batch. Any keys which could not be
-	/// fetched are elided from the resulting map.
-	pub async fn get_pubkeys<'a, S, K>(&self, batch: S) -> PubKeyMap
-	where
-		S: Iterator<Item = (&'a ServerName, K)> + Send,
-		K: Iterator<Item = &'a ServerSigningKeyId> + Send,
-	{
-		let mut keys = PubKeyMap::new();
-		for (server, key_ids) in batch {
-			let pubkeys = self.get_pubkeys_for(server, key_ids).await;
-			keys.insert(server.into(), pubkeys);
-		}
-
-		keys
-	}
-
-	/// Gets verify keys for the given server name and a set of key IDs. Any
-	/// keys that cannot be found are elided from the resulting set.
-	pub async fn get_pubkeys_for<'a, I>(&self, origin: &ServerName, key_ids: I) -> PubKeys
-	where
-		I: Iterator<Item = &'a ServerSigningKeyId> + Send,
-	{
-		let mut keys = PubKeys::new();
-		for key_id in key_ids {
-			if let Ok(verify_key) = self.get_verify_key(origin, key_id).await {
-				keys.insert(key_id.into(), verify_key.key);
-			}
-		}
-
-		keys
-	}
-
-	/// Fetches a single verify key via origin and notaries.
-	pub async fn get_verify_key(
+	/// Fetches the verify keys required to verify all the incoming events.
+	pub async fn get_events_keys(
 		&self,
-		origin: &ServerName,
-		key_id: &ServerSigningKeyId,
-	) -> Result<VerifyKey> {
-		let notary_first = self.services.server.config.query_trusted_key_servers_first;
-		let notary_only = self.services.server.config.only_query_trusted_key_servers;
-
-		if let Some(result) = self.verify_keys_for(origin).await.remove(key_id) {
-			trace!("Found key in cache");
-			return Ok(result);
-		}
-
-		if notary_first {
-			if let Ok(result) = self.get_verify_key_from_notaries(origin, key_id).await {
-				return Ok(result);
-			}
-		}
-
-		if !notary_only {
-			if let Ok(result) = self.get_verify_key_from_origin(origin, key_id).await {
-				return Ok(result);
-			}
-		}
-
-		if !notary_first {
-			if let Ok(result) = self.get_verify_key_from_notaries(origin, key_id).await {
-				return Ok(result);
-			}
-		}
-
-		Err!(BadServerResponse(debug_error!(
-			%key_id,
-			%origin,
-			"Failed to fetch federation signing-key"
-		)))
-	}
-
-	/// Fetches verify keys from notary servers, adding them to the cache, and
-	/// returning the requested key if found.
-	async fn get_verify_key_from_notaries(
-		&self,
-		origin: &ServerName,
-		key_id: &ServerSigningKeyId,
-	) -> Result<VerifyKey> {
-		for notary in self.services.globals.trusted_servers() {
-			if let Ok(server_keys) = self.notary_request(notary, origin).await {
-				for server_key in server_keys.clone() {
-					self.add_signing_keys(server_key).await;
+		raw_objects: impl Iterator<Item = &Box<RawValue>>,
+		room_version_rules: &RoomVersionRules,
+	) -> PubKeyMap {
+		let mut batch = ServerKeysQuery::new();
+		for raw_object in raw_objects {
+			let object: CanonicalJsonObject = match serde_json::from_str(raw_object.get()) {
+				| Ok(object) => object,
+				| Err(e) => {
+					debug_error!(?raw_object, error=?e, "failed to parse incoming event");
+					continue;
+				},
+			};
+			let min_valid_ts = match Self::min_valid_ts_from_event(&object) {
+				| Ok(ts) => ts,
+				| Err(e) => {
+					debug!(?object, error=?e, "incoming event has no `origin_server_ts` field, skipping");
+					continue;
+				},
+			};
+			let required = match required_keys(&object, &room_version_rules.signatures) {
+				| Ok(required) => required,
+				| Err(e) => {
+					debug_warn!(
+						?object,
+						error=?e,
+						"failed to determine required verify keys for incoming event, skipping"
+					);
+					continue;
+				},
+			};
+			for (server_name, signing_key_ids) in required {
+				let batch_entry = batch.entry(server_name).or_default();
+				for key_id in signing_key_ids {
+					batch_entry
+						.entry(key_id)
+						.and_modify(|criteria| {
+							// We want to use the lower of the two minimum valid
+							// timestamps if the key is already existing;
+							// otherwise, just insert this one.
+							criteria.minimum_valid_until_ts = Some(
+								criteria
+									.minimum_valid_until_ts
+									.expect("We never ask for no min valid ts")
+									.min(min_valid_ts),
+							);
+						})
+						.or_insert_with(|| new_query_criteria(min_valid_ts));
 				}
+			}
+		}
 
-				for server_key in server_keys {
-					if let Some(result) = extract_key(server_key, key_id) {
-						return Ok(result);
+		self.fetch_verify_keys(batch).await
+	}
+
+	/// Fetches verify keys matching the required batch parameters.
+	///
+	/// First attempts to fulfil any requirements from the local server cache,
+	/// falling back to asking origins and notaries.
+	///
+	/// Origins and notaries are always asked in parallel, and will complement
+	/// each other if notaries are not prioritised, using a
+	/// first-come-first-serve stream consumption approach.
+	///
+	/// If notaries are prioritised, this function will instead drain all notary
+	/// futures until all keys are found, or there are no more notary responses.
+	/// In that case, the origin future(s) will be consumed.
+	///
+	/// If there are still unacquired keys, a debug info log will be emitted
+	/// containing further details, but an error will not be returned.
+	async fn fetch_verify_keys(&self, batch: ServerKeysQuery) -> PubKeyMap {
+		let mut still_missing = batch.clone();
+		let mut result = PubKeyMap::new();
+		trace!(?batch, "Preparing to fetch verify keys");
+
+		// First, filter out any keys we already have locally.
+		for (origin, keys) in &batch {
+			for (key_id, criteria) in keys {
+				let min_valid_ts = criteria
+					.minimum_valid_until_ts
+					.unwrap_or_else(|| MilliSecondsSinceUnixEpoch(uint!(0)));
+
+				if let Some(server_keys) = self.get_signing_key(origin, key_id).await
+					&& let Some(key) =
+						verify_key_from_response(key_id, &server_keys, min_valid_ts)
+				{
+					trace!(%origin, %key_id, ?min_valid_ts, "Verify key exists locally");
+					Self::track_one_result(
+						&mut result,
+						&mut still_missing,
+						origin,
+						key_id,
+						key.key,
+					);
+					continue;
+				}
+				trace!(%origin, %key_id, ?min_valid_ts, "Did not find verify key locally");
+			}
+		}
+		if still_missing.is_empty() {
+			trace!("Not missing any more keys!");
+			return result;
+		}
+		trace!(?still_missing, "Still missing some keys after local population");
+
+		// Now we need to fetch keys en masse.
+		// For each server in the set, we'll ping off a request to the origin to
+		// fetch their signing keys. Then, in parallel, we'll ask every notary
+		// we know for the full batch. If the server is configured to ask the
+		// origin first, we'll simply await the origin request first, failing
+		// back to the notary futures. Otherwise, we can work through each
+		// response as they come back until we either get all the keys we
+		// need, or we run out of servers to ask.
+		let mut origin_futures: FuturesUnordered<_> = FuturesUnordered::new();
+		let mut notary_futures: FuturesUnordered<_> = FuturesUnordered::new();
+
+		for origin in still_missing.keys() {
+			let min_valid_ts_for_this_remote = batch
+				.get(origin)
+				.and_then(|keys| {
+					keys.values()
+						.filter_map(|criteria| criteria.minimum_valid_until_ts)
+						.max()
+				})
+				.unwrap_or_else(|| MilliSecondsSinceUnixEpoch(uint!(0)));
+			trace!(
+				%origin,
+				?min_valid_ts_for_this_remote,
+				"Spawning future to ask origin for their verify keys"
+			);
+			origin_futures
+				.push(self.origin_request(origin.to_owned(), min_valid_ts_for_this_remote));
+		}
+
+		self.services
+			.server
+			.config
+			.trusted_servers
+			.iter()
+			.for_each(|notary| {
+				trace!(notary=%notary.server_name(), "Asking notary for still missing keys");
+				notary_futures
+					.push(self.notary_request(notary.server_name(), still_missing.clone()));
+			});
+
+		let prioritise_notaries = !self.services.server.config.trusted_servers.is_empty()
+			&& self.services.server.config.query_trusted_key_servers_first;
+
+		// TODO: In the future, an option to establish a quorum between notaries
+		// may be desired. In such a case, the `still_missing { break }`
+		// blocks should be skipped in order to acquire all notary responses for
+		// later comparison.
+
+		if !prioritise_notaries {
+			// If we aren't prioritising notaries, we can just select futures on
+			// a first-come-first-serve basis.
+			loop {
+				select! {
+					origin_response = origin_futures.next(), if !origin_futures.is_empty() => {
+						let origin_response = match origin_response {
+							Some(Ok(v)) => v,
+							None => continue,
+							Some(Err(e)) => {
+								debug_warn!(err=?e, "Got error response from origin");
+								continue;
+							}
+						};
+						trace!(?origin_response, "Got signing key response from origin");
+						self.add_signing_keys(&origin_response);
+						Self::track_result(&mut result, &mut still_missing, &origin_response);
+						if still_missing.is_empty() {
+							break;
+						}
+					},
+					notary_response = notary_futures.next(), if !notary_futures.is_empty() => {
+						let notary_response = match notary_response {
+							Some(Ok(v)) => v,
+							None => continue,
+							Some(Err(e)) => {
+								debug_warn!(err=?e, "Got error response from notary");
+								continue;
+							}
+						};
+						trace!(
+							chunks=%notary_response.len(),
+							"Got signing keys response from a notary"
+						);
+						for server_keys in notary_response {
+							trace!(?server_keys, "Adding response from notary");
+							self.add_signing_keys(&server_keys);
+							Self::track_result(&mut result, &mut still_missing, &server_keys);
+							if still_missing.is_empty() {
+								break;
+							}
+						}
+					}
+					else => {
+						trace!("Origin or notary future stream(s) finished.");
+						break
+					}
+				}
+			}
+		} else {
+			// Drain notary responses first, and then sift through origins.
+			trace!(count=%notary_futures.len(), "Waiting for notary responses...");
+			while let Some(res) = notary_futures.next().await {
+				let Ok(res) = res else {
+					debug_warn!(err=?res.err(), "Got error response from notary");
+					continue;
+				};
+				for server_keys in res {
+					trace!(?server_keys, "Adding response from notary");
+					self.add_signing_keys(&server_keys);
+					Self::track_result(&mut result, &mut still_missing, &server_keys);
+				}
+				if still_missing.is_empty() {
+					break;
+				}
+			}
+			if !still_missing.is_empty() {
+				trace!("Still missing some keys, waiting for origin responses...");
+				while let Some(res) = origin_futures.next().await {
+					let Ok(res) = res else {
+						debug_warn!(err=?res.err(), "Got error response from origin");
+						continue;
+					};
+					trace!(response=?res, "Got signing key response from origin");
+					self.add_signing_keys(&res);
+					Self::track_result(&mut result, &mut still_missing, &res);
+					if still_missing.is_empty() {
+						break;
 					}
 				}
 			}
 		}
 
-		Err!(Request(NotFound("Failed to fetch signing-key from notaries")))
+		if !still_missing.is_empty() {
+			let missing_servers = still_missing.len();
+			let missing_keys = still_missing
+				.values()
+				.map(BTreeMap::len)
+				.fold(0_usize, usize::saturating_add);
+			let fetched_servers = result.len();
+			let fetched_keys = result
+				.values()
+				.map(BTreeMap::len)
+				.fold(0_usize, usize::saturating_add);
+			debug_info!(
+				%fetched_servers,
+				%fetched_keys,
+				%missing_servers,
+				%missing_keys,
+				"Could not acquire all requested signing keys"
+			);
+			trace!(?still_missing, "Could not acquire all requested signing keys");
+		}
+		result
 	}
 
-	/// Fetches verify keys from the origin server, adding them to the cache,
-	/// and returning the requested key if found.
-	async fn get_verify_key_from_origin(
+	/// Tracks the result of a signing keys response by removing each of the
+	/// returned verify keys and old verify keys from the missing map, and
+	/// adding them to the results map.
+	fn track_result(
+		results: &mut PubKeyMap,
+		missing: &mut ServerKeysQuery,
+		batch: &ServerSigningKeys,
+	) {
+		batch.verify_keys.iter().for_each(|(key_id, key)| {
+			Self::track_one_result(results, missing, &batch.server_name, key_id, key.key.clone());
+		});
+		batch.old_verify_keys.iter().for_each(|(key_id, key)| {
+			Self::track_one_result(results, missing, &batch.server_name, key_id, key.key.clone());
+		});
+	}
+
+	/// Tracks a single signing key result by adding it to the results map and
+	/// removing it from the missing map. If all missing keys for an origin are
+	/// found (the missing list is empty), the origin is removed from the
+	/// missing map entirely.
+	fn track_one_result(
+		results: &mut PubKeyMap,
+		missing: &mut ServerKeysQuery,
+		origin: &ServerName,
+		key_id: &ServerSigningKeyId,
+		key: Base64,
+	) {
+		results
+			.entry(origin.to_string())
+			.or_default()
+			.insert(key_id.to_string(), key);
+		missing.entry(origin.to_owned()).and_modify(|keys| {
+			keys.remove(key_id);
+		});
+		if missing.get(origin).is_some_and(BTreeMap::is_empty) {
+			missing.remove(origin);
+		}
+	}
+
+	/// Fetches a single verify key.
+	///
+	/// If a valid key cannot be found locally, the origin is queried. If it
+	/// fails to respond, or the desired verify key is not contained in the
+	/// response/expired before `min_valid_ts`, `None` is returned.
+	#[tracing::instrument(skip(self))]
+	pub async fn get_single_verify_key(
 		&self,
 		origin: &ServerName,
 		key_id: &ServerSigningKeyId,
-	) -> Result<VerifyKey> {
-		if let Ok(server_key) = self.server_request(origin).await {
-			self.add_signing_keys(server_key.clone()).await;
-			if let Some(result) = extract_key(server_key, key_id) {
-				return Ok(result);
-			}
+		min_valid_ts: MilliSecondsSinceUnixEpoch,
+	) -> Option<VerifyKey> {
+		if let Some(vk) = self
+			.get_signing_key(origin, key_id)
+			.await
+			.and_then(|ssk| verify_key_from_response(key_id, &ssk, min_valid_ts))
+		{
+			trace!("Found verify key locally");
+			return Some(vk);
 		}
 
-		Err!(Request(NotFound("Failed to fetch signing-key from origin")))
+		trace!("Local miss for single key, asking origin");
+		let response = self
+			.origin_request(origin.to_owned(), min_valid_ts)
+			.await
+			.inspect_err(|e| info!("Could not acquire signing key from origin server: {e:?}"))
+			.ok()?;
+
+		verify_key_from_response(key_id, &response, min_valid_ts)
+			.inspect(|_| self.add_signing_keys(&response))
+	}
+}
+
+/// Generates a QueryCriteria with the minimum_valid_until_ts set to the target
+/// epoch.
+fn new_query_criteria(min_valid_ts: MilliSecondsSinceUnixEpoch) -> QueryCriteria {
+	assign!(QueryCriteria::new(), { minimum_valid_until_ts: Some(min_valid_ts) })
+}
+
+/// Extracts a verify key with the given ID from the response.
+///
+/// If the provided `ServerSigningKeys` is dated to before the min valid TS,
+/// this will throw an assertion panic in debug builds, or always return `None`
+/// in release. If the key is found in `verify_keys`, it is returned as-is. If
+/// the key is found in `old_verify_keys`, it is only returned if the expiry
+/// timestamp is greater than or equal to the minimum valid timestamp.
+fn verify_key_from_response(
+	key_id: &ServerSigningKeyId,
+	server_keys: &ServerSigningKeys,
+	min_valid_ts: MilliSecondsSinceUnixEpoch,
+) -> Option<VerifyKey> {
+	if server_keys.valid_until_ts < min_valid_ts || server_keys.valid_until_ts > in_one_week() {
+		return None;
+	}
+	if let Some(key) = server_keys
+		.verify_keys
+		.iter()
+		.filter(|(k, _)| *k == key_id)
+		.map(|(_, v)| v)
+		.next()
+	{
+		return Some(key.to_owned());
+	}
+	server_keys
+		.old_verify_keys
+		.iter()
+		.filter(|(k, v)| {
+			*k == key_id && v.expired_ts.min(server_keys.valid_until_ts) >= min_valid_ts
+		})
+		.map(|(_, v)| v)
+		.next()
+		.map(|old_key| VerifyKey::new(old_key.key.clone()))
+}
+
+#[cfg(test)]
+mod test {
+	use ruma::Int;
+	use serde_json::json;
+
+	use super::*;
+	use crate::server_keys::{Service, request::ServerKeysQuerySet};
+
+	#[test]
+	fn min_valid_ts_refuses_garbage() {
+		let garbage = vec![
+			json!("string"),
+			json!("0"), // string ints aren't ints
+			CanonicalJsonValue::Null.into(),
+			json!([]),
+			json!({}),
+			json!(" "),
+		];
+
+		for fixture in garbage {
+			let obj: CanonicalJsonObject =
+				serde_json::from_value(json!({"origin_server_ts": fixture})).unwrap();
+			assert!(
+				Service::min_valid_ts_from_event(&obj).is_err(),
+				"garbage should always error instead of trying to guess"
+			);
+		}
+	}
+
+	#[test]
+	fn min_valid_ts_requires_field() {
+		// This just ensures we don't accidentally default to 0 or something
+		// dumb like that
+		assert!(
+			Service::min_valid_ts_from_event(&CanonicalJsonObject::new()).is_err(),
+			"origin_server_ts should be a required field and thus should error if not provided"
+		);
+	}
+
+	#[test]
+	fn min_valid_ts_accepts_negative_values() {
+		// Since we convert to u64 with saturating for MSSinceEpoch, this should
+		// give us a final timestamp of zero (since unsigned ints cant be
+		// negative).
+		let obj: CanonicalJsonObject =
+			serde_json::from_value(json!({"origin_server_ts": Int::MIN})).unwrap();
+		assert_eq!(
+			Service::min_valid_ts_from_event(&obj)
+				.expect("Int::MIN should be a valid min timestamp")
+				.0,
+			uint!(0),
+			"Saturating Int::MIN to u64 should result in zero"
+		);
+	}
+
+	#[test]
+	fn tracking_server_result_removes_empty_server_from_missing_map() {
+		let server_name = ruma::owned_server_name!("example.com");
+		let key_id = ServerSigningKeyId::parse("ed25519:1").unwrap();
+		let mut missing = ServerKeysQuery::new();
+		missing.insert(server_name.clone(), ServerKeysQuerySet::new());
+		missing.entry(server_name.clone()).and_modify(|set| {
+			set.insert(key_id.clone(), QueryCriteria::new());
+		});
+		let mut results = PubKeyMap::new();
+
+		Service::track_one_result(
+			&mut results,
+			&mut missing,
+			&server_name,
+			&key_id,
+			Base64::new(Vec::new()),
+		);
+
+		assert!(missing.is_empty(), "missing should be empty after tracking the result");
+		assert!(!results.is_empty(), "results should not be empty after tracking the result");
 	}
 }

src/service/server_keys/keypair.rs

@@ -9,7 +9,8 @@ use super::VerifyKeys;
 pub(super) fn init(db: &Arc<Database>) -> Result<(Box<Ed25519KeyPair>, VerifyKeys)> {
 	let keypair = load(db).inspect_err(|_e| {
 		error!("Keypair invalid. Deleting...");
-		remove(db);
+		let global = &db["global"];
+		global.remove(b"keypair");
 	})?;
 
 	let verify_key = VerifyKey::new(Base64::new(keypair.public_key().to_vec()));
@@ -20,6 +21,7 @@ pub(super) fn init(db: &Arc<Database>) -> Result<(Box<Ed25519KeyPair>, VerifyKey
 	Ok((keypair, verify_keys))
 }
 
+/// Loads the existing keypair from the database, otherwise creates a new one.
 fn load(db: &Arc<Database>) -> Result<Box<Ed25519KeyPair>> {
 	let (version, key) = db["global"]
 		.get_blocking(b"keypair")
@@ -44,6 +46,7 @@ fn load(db: &Arc<Database>) -> Result<Box<Ed25519KeyPair>> {
 	Ok(Box::new(key))
 }
 
+/// Creates a new keypair and writes it directly to the db.
 fn create(db: &Arc<Database>) -> Result<(String, Vec<u8>)> {
 	let keypair = Ed25519KeyPair::generate();
 
@@ -55,9 +58,3 @@ fn create(db: &Arc<Database>) -> Result<(String, Vec<u8>)> {
 
 	Ok(value)
 }
-
-#[inline]
-fn remove(db: &Arc<Database>) {
-	let global = &db["global"];
-	global.remove(b"keypair");
-}

src/service/server_keys/mod.rs

@@ -1,4 +1,3 @@
-mod acquire;
 mod get;
 mod keypair;
 mod request;
@@ -6,30 +5,28 @@ mod sign;
 mod util;
 mod verify;
 
-use std::{collections::BTreeMap, sync::Arc, time::Duration};
+use std::{collections::BTreeMap, sync::Arc};
 
 use conduwuit::{
 	Result, Server,
-	utils::{IterStream, timepoint_from_now},
+	utils::{IterStream, ReadyExt, stream::TryIgnore},
 };
-use database::{Deserialized, Json, Map};
-use futures::StreamExt;
+use database::{Deserialized, Ignore, Interfix, Json, Map};
+use futures::{Stream, StreamExt};
+pub use request::in_one_week;
 use ruma::{
 	CanonicalJsonObject, MilliSecondsSinceUnixEpoch, OwnedServerSigningKeyId, ServerName,
 	ServerSigningKeyId,
 	api::federation::discovery::{ServerSigningKeys, VerifyKey},
 	room_version_rules::RoomVersionRules,
-	serde::Raw,
 	signatures::{Ed25519KeyPair, PublicKeyMap, PublicKeySet},
 };
-use serde_json::value::RawValue as RawJsonValue;
 
 use crate::{Dep, globals, sending, server_keys::util::required_keys};
 
 pub struct Service {
 	keypair: Box<Ed25519KeyPair>,
 	verify_keys: VerifyKeys,
-	minimum_valid: Duration,
 	services: Services,
 	db: Data,
 }
@@ -41,7 +38,7 @@ struct Services {
 }
 
 struct Data {
-	server_signingkeys: Arc<Map>,
+	servernamekeyid_response: Arc<Map>,
 }
 
 pub type VerifyKeys = BTreeMap<OwnedServerSigningKeyId, VerifyKey>;
@@ -50,22 +47,19 @@ pub type PubKeys = PublicKeySet;
 
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
-		let minimum_valid = Duration::from_hours(1);
-
 		let (keypair, verify_keys) = keypair::init(args.db)?;
 		debug_assert!(verify_keys.len() == 1, "only one active verify_key supported");
 
 		Ok(Arc::new(Self {
 			keypair,
 			verify_keys,
-			minimum_valid,
 			services: Services {
 				globals: args.depend::<globals::Service>("globals"),
 				sending: args.depend::<sending::Service>("sending"),
 				server: args.server.clone(),
 			},
 			db: Data {
-				server_signingkeys: args.db["server_signingkeys"].clone(),
+				servernamekeyid_response: args.db["servernamekeyid_response"].clone(),
 			},
 		}))
 	}
@@ -77,40 +71,32 @@ impl Service {
 	/// Fetches the server's active keypair.
 	pub fn keypair(&self) -> &Ed25519KeyPair { &self.keypair }
 
-	/// Fetches the server's active signing key ID.
-	pub fn active_key_id(&self) -> &ServerSigningKeyId { self.active_verify_key().0 }
-
 	/// Fetches the server's active signing key. Panics if there's more than 1
 	/// active key.
-	pub fn active_verify_key(&self) -> (&ServerSigningKeyId, &VerifyKey) {
+	pub fn active_verify_key(&self) -> (OwnedServerSigningKeyId, VerifyKey) {
 		debug_assert!(self.verify_keys.len() <= 1, "more than one active verify_key");
 		self.verify_keys
 			.iter()
 			.next()
-			.map(|(id, key)| (id.as_ref(), key))
+			.map(|(id, key)| (id.to_owned(), key.to_owned()))
 			.expect("missing active verify_key")
 	}
 
-	/// Combines signing keys from the incoming response to the existing
-	/// database set.
-	async fn add_signing_keys(&self, new_keys: ServerSigningKeys) {
-		let origin = &new_keys.server_name;
-
-		// (timo) Not atomic, but this is not critical
-		let mut keys: ServerSigningKeys = self
-			.db
-			.server_signingkeys
-			.get(origin)
-			.await
-			.deserialized()
-			.unwrap_or_else(|_| {
-				// Just insert "now", it doesn't matter
-				ServerSigningKeys::new(origin.to_owned(), MilliSecondsSinceUnixEpoch::now())
-			});
-
-		keys.verify_keys.extend(new_keys.verify_keys);
-		keys.old_verify_keys.extend(new_keys.old_verify_keys);
-		self.db.server_signingkeys.raw_put(origin, Json(&keys));
+	/// Adds a signing key response to the database.
+	fn add_signing_keys(&self, new_keys: &ServerSigningKeys) {
+		// TODO(nex): this results in massive amounts of data duplication.
+		// Perhaps a sliding approach would be better? where we just update the
+		// last known response if the only thing that changed was the
+		// valid_until_ts. This would require a more complex data structure
+		// though. For now, just storing the duplicated data is fine.
+		for key in new_keys
+			.verify_keys
+			.keys()
+			.chain(new_keys.old_verify_keys.keys())
+		{
+			let db_key = (new_keys.server_name.as_str(), key.as_str());
+			self.db.servernamekeyid_response.put(db_key, Json(new_keys));
+		}
 	}
 
 	/// Checks if the keys required to verify an incoming PDU are already known
@@ -123,11 +109,14 @@ impl Service {
 		let Ok(required_keys) = required_keys(object, &room_version_rules.signatures) else {
 			return false;
 		};
+		let Ok(min_valid_ts) = Self::min_valid_ts_from_event(object) else {
+			return false;
+		};
 		required_keys
 			.iter()
 			.flat_map(|(server, key_ids)| key_ids.iter().map(move |key_id| (server, key_id)))
 			.stream()
-			.all(|(server, key_id)| self.verify_key_exists(server, key_id))
+			.all(|(server, key_id)| self.verify_key_exists(server, key_id, min_valid_ts))
 			.await
 	}
 
@@ -137,40 +126,54 @@ impl Service {
 		&self,
 		origin: &ServerName,
 		key_id: &ServerSigningKeyId,
+		minimum_valid_ts: MilliSecondsSinceUnixEpoch,
 	) -> bool {
-		type KeysMap = BTreeMap<OwnedServerSigningKeyId, Box<RawJsonValue>>;
-
-		let Ok(keys) = self
-			.db
-			.server_signingkeys
-			.get(origin)
-			.await
-			.deserialized::<Raw<ServerSigningKeys>>()
-		else {
+		let Some(keys) = self.get_signing_key(origin, key_id).await else {
+			// We have no keys for this server
 			return false;
 		};
 
-		if let Ok(Some(verify_keys)) = keys.get_field::<KeysMap>("verify_keys") {
-			if verify_keys.contains_key(key_id) {
-				return true;
-			}
+		if keys.valid_until_ts < minimum_valid_ts || keys.valid_until_ts > in_one_week() {
+			// This key response is stale and cannot be trusted.
+			return false;
 		}
 
-		if let Ok(Some(old_verify_keys)) = keys.get_field::<KeysMap>("old_verify_keys") {
-			if old_verify_keys.contains_key(key_id) {
-				return true;
-			}
+		if keys.verify_keys.contains_key(key_id) {
+			// An in-date, active key exists for this server.
+			return true;
+		}
+
+		if let Some(expired_key) = keys.old_verify_keys.get(key_id) {
+			// This key can still be used to verify signatures provided it
+			// expired within the minimum valid time window.
+			return expired_key.expired_ts.min(keys.valid_until_ts) >= minimum_valid_ts;
 		}
 
 		false
 	}
 
-	/// Returns all verify keys for the origin.
+	pub async fn get_signing_key(
+		&self,
+		origin: &ServerName,
+		key_id: &ServerSigningKeyId,
+	) -> Option<ServerSigningKeys> {
+		self.db
+			.servernamekeyid_response
+			.qry(&(origin.as_str(), key_id.as_str()))
+			.await
+			.deserialized::<ServerSigningKeys>()
+			.ok()
+	}
+
+	/// Returns all (active) verify keys for the origin.
 	pub async fn verify_keys_for(&self, origin: &ServerName) -> VerifyKeys {
 		let mut keys = self
 			.signing_keys_for(origin)
-			.await
-			.map_or(BTreeMap::new(), |keys| merge_old_keys(keys).verify_keys);
+			.ready_fold(BTreeMap::new(), |mut base, keys| {
+				base.extend(keys.verify_keys);
+				base
+			})
+			.await;
 
 		if self.services.globals.server_is_ours(origin) {
 			keys.extend(self.verify_keys.clone());
@@ -179,39 +182,15 @@ impl Service {
 		keys
 	}
 
-	/// Returns the stored server signing keys response for the origin.
-	pub async fn signing_keys_for(&self, origin: &ServerName) -> Result<ServerSigningKeys> {
-		self.db.server_signingkeys.get(origin).await.deserialized()
-	}
-
-	/// Returns the minimum_valid_ts for a key response from the current point
-	/// in time, plus the minimum valid time.
-	fn minimum_valid_ts(&self) -> MilliSecondsSinceUnixEpoch {
-		let timepoint =
-			timepoint_from_now(self.minimum_valid).expect("SystemTime should not overflow");
-		MilliSecondsSinceUnixEpoch::from_system_time(timepoint).expect("UInt should not overflow")
+	/// Returns the stored server signing keys responses for the origin.
+	pub fn signing_keys_for<'a>(
+		&'a self,
+		origin: &'a ServerName,
+	) -> impl Stream<Item = ServerSigningKeys> + Send + 'a {
+		self.db
+			.servernamekeyid_response
+			.stream_prefix(&(origin, Interfix))
+			.ignore_err()
+			.map(|(_, v): (Ignore, ServerSigningKeys)| v)
 	}
 }
-
-fn merge_old_keys(mut keys: ServerSigningKeys) -> ServerSigningKeys {
-	keys.verify_keys.extend(
-		keys.old_verify_keys
-			.clone()
-			.into_iter()
-			.map(|(key_id, old)| (key_id, VerifyKey::new(old.key))),
-	);
-
-	keys
-}
-
-fn extract_key(mut keys: ServerSigningKeys, key_id: &ServerSigningKeyId) -> Option<VerifyKey> {
-	keys.verify_keys.remove(key_id).or_else(|| {
-		keys.old_verify_keys
-			.remove(key_id)
-			.map(|old| VerifyKey::new(old.key))
-	})
-}
-
-fn key_exists(keys: &ServerSigningKeys, key_id: &ServerSigningKeyId) -> bool {
-	keys.verify_keys.contains_key(key_id) || keys.old_verify_keys.contains_key(key_id)
-}

src/service/server_keys/request.rs

@@ -1,110 +1,39 @@
-use std::{collections::BTreeMap, fmt::Debug};
+use std::collections::BTreeMap;
 
-use conduwuit::{Err, Result, debug};
-use get_remote_server_keys_batch::v2::Request;
+use conduwuit::{Err, Result, trace, utils::millis_since_unix_epoch, warn};
 use ruma::{
-	OwnedServerName, OwnedServerSigningKeyId, ServerName, ServerSigningKeyId,
+	MilliSecondsSinceUnixEpoch, OwnedServerName, OwnedServerSigningKeyId, ServerName, UInt,
 	api::federation::discovery::{
-		ServerSigningKeys, get_remote_server_keys,
-		get_remote_server_keys_batch::{self, v2::QueryCriteria},
+		ServerSigningKeys,
+		get_remote_server_keys_batch::{
+			self,
+			v2::{QueryCriteria, Response as NotaryResponse},
+		},
 		get_server_keys,
 	},
+	serde::Base64,
+	uint,
 };
-type RumaBatch = BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>;
+
+pub(crate) type ServerKeysQuery = BTreeMap<OwnedServerName, ServerKeysQuerySet>;
+pub(crate) type ServerKeysQuerySet = BTreeMap<OwnedServerSigningKeyId, QueryCriteria>;
 
 impl super::Service {
-	/// Handles a batched notary request.
-	/// If the batch is empty, an empty vec is returned.
-	pub(super) async fn batch_notary_request<'a, S, K>(
+	/// Asks the origin server directly for its signing keys.
+	///
+	/// An error is returned if the server responds with invalid keys; this is
+	/// in contrast with the notary request, which simply drops such responses.
+	pub async fn origin_request(
 		&self,
-		notary: &ServerName,
-		batch: S,
-	) -> Result<Vec<ServerSigningKeys>>
-	where
-		S: Iterator<Item = (&'a ServerName, K)> + Send,
-		K: Iterator<Item = &'a ServerSigningKeyId> + Send,
-	{
-		let mut criteria = QueryCriteria::new();
-		criteria.minimum_valid_until_ts = Some(self.minimum_valid_ts());
-
-		let mut server_keys = batch.fold(RumaBatch::new(), |mut batch, (server, key_ids)| {
-			batch
-				.entry(server.into())
-				.or_default()
-				.extend(key_ids.map(|key_id| (key_id.into(), criteria.clone())));
-
-			batch
-		});
-
-		if server_keys.is_empty() {
-			return Ok(vec![]);
-		}
-
-		let mut results = Vec::new();
-		while let Some(batch) = server_keys
-			.keys()
-			.rev()
-			.take(self.services.server.config.trusted_server_batch_size)
-			.next_back()
-			.cloned()
-		{
-			let request = Request::new(server_keys.split_off(&batch));
-
-			debug!(
-				?notary,
-				?batch,
-				remaining = %server_keys.len(),
-				requesting = ?request.server_keys.keys(),
-				"notary request"
-			);
-
-			let response = self
-				.services
-				.sending
-				.send_unauthenticated_request(notary, request)
-				.await?
-				.server_keys
-				.into_iter()
-				.map(|key| key.deserialize())
-				.filter_map(Result::ok);
-
-			results.extend(response);
-		}
-
-		Ok(results)
-	}
-
-	/// Makes a request to a notary server to fetch signing keys.
-	pub async fn notary_request(
-		&self,
-		notary: &ServerName,
-		target: &ServerName,
-	) -> Result<impl Iterator<Item = ServerSigningKeys> + Clone + Debug + Send + use<>> {
-		use get_remote_server_keys::v2::Request;
-
-		let request = Request::new(target.into(), self.minimum_valid_ts());
-
-		let response = self
-			.services
-			.sending
-			.send_unauthenticated_request(notary, request)
-			.await?
-			.server_keys
-			.into_iter()
-			.map(|key| key.deserialize())
-			.filter_map(Result::ok);
-
-		Ok(response)
-	}
-
-	/// Makes a request to an origin server to fetch signing keys.
-	pub async fn server_request(&self, target: &ServerName) -> Result<ServerSigningKeys> {
+		target: OwnedServerName,
+		min_valid_ts: MilliSecondsSinceUnixEpoch,
+	) -> Result<ServerSigningKeys> {
 		use get_server_keys::v2::Request;
 
 		let server_signing_key = self
 			.services
 			.sending
-			.send_unauthenticated_request(target, Request::new())
+			.send_unauthenticated_request(&target, Request::new())
 			.await
 			.map(|response| response.server_key)
 			.and_then(|key| key.deserialize().map_err(Into::into))?;
@@ -113,10 +42,172 @@ impl super::Service {
 			return Err!(BadServerResponse(debug_warn!(
 				requested = ?target,
 				response = ?server_signing_key.server_name,
-				"Server responded with bogus server_name"
+				"Origin server responded with bogus server_name"
+			)));
+		}
+
+		if server_signing_key.valid_until_ts < min_valid_ts
+			|| server_signing_key.valid_until_ts > in_one_week()
+		{
+			return Err!(BadServerResponse(debug_warn!(
+				?target,
+				?min_valid_ts,
+				?server_signing_key.valid_until_ts,
+				"Origin server responded with server keys that are not in-date"
+			)));
+		}
+
+		if let Err(e) = Self::verify_server_keys_response(&server_signing_key, None) {
+			return Err!(BadServerResponse(debug_warn!(
+				?target,
+				"Origin server failed to self-sign key response: {e}"
 			)));
 		}
 
 		Ok(server_signing_key)
 	}
+
+	/// Sends a request to a notary server, asking for the specified keys for
+	/// the given servers.
+	///
+	/// The response is filtered before it is returned, dropping any chunks that
+	/// are invalid or do not match the requested keys. As a result, the request
+	/// may be successful, but the returned vec may be empty.
+	///
+	///
+	/// An error is returned if the notary server itself returns an error,
+	/// because of invalid data.
+	///
+	/// Does not make a request if there are no keys specified in the batch.
+	pub async fn notary_request(
+		&self,
+		notary: &ServerName,
+		batch: ServerKeysQuery,
+	) -> Result<Vec<ServerSigningKeys>> {
+		use get_remote_server_keys_batch::v2::Request;
+
+		let is_empty = batch.is_empty() || batch.values().all(BTreeMap::is_empty);
+		if is_empty {
+			return Ok(Vec::new());
+		}
+
+		let request = Request::new(batch.clone());
+		let notary_keys = self
+			.services
+			.server
+			.config
+			.trusted_servers
+			.iter()
+			.find(|n| n.server_name() == notary)
+			.map(|nk| nk.verify_keys().to_owned())
+			.unwrap_or_default();
+		let keys = self
+			.services
+			.sending
+			.send_unauthenticated_request(notary, request)
+			.await
+			.map(|response| Self::parse_notary_response(&response))?
+			.into_iter()
+			.filter(|keys| {
+				Self::validate_notary_chunk(&batch, keys, notary, notary_keys.as_slice())
+			})
+			.collect();
+		Ok(keys)
+	}
+
+	/// Parses a notary response, returning a vec of server signing key
+	/// responses. Drops any chunks that cannot be deserialized.
+	fn parse_notary_response(raw: &NotaryResponse) -> Vec<ServerSigningKeys> {
+		raw.server_keys
+			.iter()
+			.map(ruma::serde::Raw::deserialize)
+			.filter_map(Result::ok)
+			.collect()
+	}
+
+	/// Validates a signing keys response received from a notary. Returns false
+	/// if the response should be dropped, true if it should be kept.
+	fn validate_notary_chunk(
+		batch: &ServerKeysQuery,
+		keys: &ServerSigningKeys,
+		notary_name: &ServerName,
+		notary_keys: &[Base64],
+	) -> bool {
+		let Some(expected_keys) = batch.get(&keys.server_name) else {
+			trace!("Dropping signing keys chunk for unrequested server: {}", keys.server_name);
+			return false;
+		};
+		let min_valid_ts = expected_keys
+			.values()
+			.map(|q| {
+				q.minimum_valid_until_ts
+					.unwrap_or_else(|| MilliSecondsSinceUnixEpoch(uint!(0)))
+			})
+			.min()
+			.unwrap_or_else(|| MilliSecondsSinceUnixEpoch(uint!(0)));
+
+		// If this notary response is older than any of the min_valid_ts we
+		// requested, then we should drop it.
+		if keys.valid_until_ts < min_valid_ts || keys.valid_until_ts > in_one_week() {
+			trace!(
+				?min_valid_ts,
+				?keys.valid_until_ts,
+				"Dropping signing keys chunk for server {} because it is older than the minimum \
+				 valid timestamp we requested",
+				keys.server_name
+			);
+			return false;
+		}
+
+		// Does this chunk even contain any keys we asked for?
+		// Note that this enforces min_valid_ts per-key, instead of using the
+		// lowest denominator.
+		let all_key_ids = keys
+			.verify_keys
+			.keys()
+			.chain(keys.old_verify_keys.keys())
+			.map(ToOwned::to_owned)
+			.collect::<Vec<_>>();
+		let interesting = expected_keys.iter().any(|(k, q)| {
+			all_key_ids.contains(k)
+				&& q.minimum_valid_until_ts
+					.is_none_or(|m| m <= keys.valid_until_ts)
+		});
+		if !interesting {
+			trace!(
+				"Dropping signing keys chunk for server {} because it does not contain any keys \
+				 we requested",
+				keys.server_name
+			);
+			return false;
+		}
+
+		Self::verify_server_keys_response(keys, Some((notary_name, notary_keys)))
+			.inspect_err(|e| {
+				trace!(
+					?e,
+					?keys,
+					?notary_name,
+					?notary_keys,
+					"Dropping illegal server signing keys response chunk"
+				);
+				warn!(
+					%notary_name,
+					server_name=%keys.server_name,
+					"Dropping illegal server signing keys response chunk: {e:?}"
+				);
+			})
+			.is_ok()
+	}
 }
+
+#[must_use]
+pub fn in_one_week() -> MilliSecondsSinceUnixEpoch {
+	MilliSecondsSinceUnixEpoch(
+		UInt::new_saturating(millis_since_unix_epoch()).saturating_add(uint!(604_800_000)),
+	)
+}
+
+// TODO(nex): Unit tests for validate_notary_chunk would be really good, but are
+// hard to do by nature of time advancing, signatures, and external notaries.
+// We can rely on integration tests for now.

src/service/server_keys/util.rs

@@ -7,7 +7,10 @@ use ruma::{
 	signatures::{VerificationError, required_server_signatures_to_verify_event},
 };
 
-/// Extracts the server names and key ids to check signatures for given event.
+/// Extracts the server names and key ids required to check signatures for the
+/// given event.
+///
+/// Spec: https://spec.matrix.org/v1.19/server-server-api/#validating-hashes-and-signatures-on-received-events
 pub(super) fn required_keys(
 	object: &CanonicalJsonObject,
 	rules: &SignaturesRules,

src/service/server_keys/verify.rs

@@ -1,44 +1,32 @@
-use conduwuit::{Err, Result, matrix::event::gen_event_id_canonical_json};
+use conduwuit::{
+	Err, Result, matrix::event::gen_event_id_canonical_json, trace, utils::to_canonical_object,
+};
 use ruma::{
-	CanonicalJsonObject, CanonicalJsonValue, OwnedEventId, room_version_rules::RoomVersionRules,
+	CanonicalJsonObject, CanonicalJsonValue, OwnedEventId, ServerName, SigningKeyAlgorithm,
+	api::federation::discovery::ServerSigningKeys,
+	room_version_rules::RoomVersionRules,
+	serde::{Base64, base64::Standard},
 	signatures::Verified,
 };
 use serde_json::value::RawValue as RawJsonValue;
 
 impl super::Service {
-	/// Validates the incoming event, and then inserts the calculated event ID
-	/// into the event_id field.
-	pub async fn validate_and_add_event_id(
-		&self,
-		pdu: &RawJsonValue,
-		room_version_rules: &RoomVersionRules,
-	) -> Result<(OwnedEventId, CanonicalJsonObject)> {
-		let (event_id, mut value) = gen_event_id_canonical_json(pdu, room_version_rules)?;
-		if let Err(e) = self.verify_event(&value, room_version_rules).await {
-			return Err!(BadServerResponse(debug_error!(
-				"Event {event_id} failed verification: {e:?}"
-			)));
-		}
-
-		value.insert("event_id".into(), CanonicalJsonValue::String(event_id.as_str().into()));
-
-		Ok((event_id, value))
-	}
-
 	/// Validates the incoming event only using locally known verification keys.
 	/// Inserts the calculated event ID into the event_id field upon success.
-	pub async fn validate_and_add_event_id_no_fetch(
+	pub async fn verify_event_json_no_fetch_add_event_id(
 		&self,
 		pdu: &RawJsonValue,
 		room_version_rules: &RoomVersionRules,
 	) -> Result<(OwnedEventId, CanonicalJsonObject)> {
+		// TODO(nex): this function is used in two places (same file) and is
+		// awfully named, can we just kill it?
 		let (event_id, mut value) = gen_event_id_canonical_json(pdu, room_version_rules)?;
 		if !self.required_keys_exist(&value, room_version_rules).await {
 			return Err!(BadServerResponse(debug_warn!(
-				"Event {event_id} cannot be verified: missing keys."
+				"Event {event_id} cannot be verified: missing verify key(s) locally."
 			)));
 		}
-		if let Err(e) = self.verify_event(&value, room_version_rules).await {
+		if let Err(e) = self.verify_event_json(&value, room_version_rules).await {
 			return Err!(BadServerResponse(debug_error!(
 				"Event {event_id} failed verification: {e:?}"
 			)));
@@ -50,7 +38,7 @@ impl super::Service {
 	}
 
 	/// Verifies the incoming event after fetching the keys required to do so.
-	pub async fn verify_event(
+	pub async fn verify_event_json(
 		&self,
 		event: &CanonicalJsonObject,
 		room_version_rules: &RoomVersionRules,
@@ -59,15 +47,267 @@ impl super::Service {
 		ruma::signatures::verify_event(&keys, event, room_version_rules).map_err(Into::into)
 	}
 
-	/// Verifies an arbitrary JSON object after fetching the keys required to do
-	/// so.
-	pub async fn verify_json(
-		&self,
-		event: &CanonicalJsonObject,
-		room_version_rules: &RoomVersionRules,
-	) -> Result {
-		// TODO: is this actually used anywhere?
-		let keys = self.get_event_keys(event, room_version_rules).await?;
-		ruma::signatures::verify_json(&keys, event).map_err(Into::into)
+	/// Verifies a server signing keys response, either as the
+	/// result of `GET /_matrix/key/v2/server`, or a chunk from
+	/// `POST /_matrix/key/v2/query`.
+	///
+	/// Returns an error if the self-signature is invalid, or strict notary
+	/// validation is used (`notary` is non-empty) and the notary did not sign
+	/// the response.
+	///
+	/// `notary` must either be empty or have ONE entry. Extraneous notaries
+	/// will be ignored.
+	pub fn verify_server_keys_response(
+		server_keys: &ServerSigningKeys,
+		notary: Option<(&ServerName, &[Base64])>,
+	) -> Result<()> {
+		if let Some((notary_name, notary_keys)) = notary {
+			if !notary_keys.is_empty() {
+				Self::verify_notary_signature(server_keys, notary_name, notary_keys)?;
+			}
+		}
+
+		// NOTE: old_verify_keys aren't considered here because old keys can't
+		// be used to sign key responses.
+		let mut set = ruma::signatures::PublicKeySet::new();
+		server_keys.verify_keys.iter().for_each(|(key_id, key)| {
+			set.insert(key_id.to_string(), key.key.clone());
+		});
+
+		let mut pubkey_map = ruma::signatures::PublicKeyMap::new();
+		pubkey_map.insert(server_keys.server_name.to_string(), set);
+
+		let mut canonical = to_canonical_object(server_keys)?;
+		Self::strip_extraneous_signatures(&mut canonical, &server_keys.server_name, &[]);
+		ruma::signatures::verify_json(&pubkey_map, &canonical).map_err(Into::into)
+	}
+
+	/// Verifies that the notary has at least one valid signature present with
+	/// the provided signing keys.
+	fn verify_notary_signature(
+		server_keys: &ServerSigningKeys,
+		notary_name: &ServerName,
+		notary_keys: &[Base64],
+	) -> Result<()> {
+		// This strictly verifies that the notary signed this response.
+		// This is used to mitigate exploits where a notary's TLS
+		// certificate has been compromised, which would otherwise allow
+		// them to produce attacker-controlled signing keys with valid
+		// self-signatures.
+		let Some(notary_signatures) = server_keys.signatures.get(notary_name) else {
+			return Err!(Request(Forbidden(
+				"Missing notary server signature from {notary_name}"
+			)));
+		};
+		let mut canonical_object = to_canonical_object(notary_signatures)?;
+		Self::strip_extraneous_signatures(&mut canonical_object, &server_keys.server_name, &[
+			notary_name,
+		]);
+		let for_verify = ruma::signatures::to_canonical_json_string_for_signing(
+			&to_canonical_object(server_keys).expect("server keys object must be canonical json"),
+		)
+		.expect("canonical JSON must be stringable");
+
+		let raw_notary_keys = notary_keys.iter().map(Base64::as_bytes).collect::<Vec<_>>();
+		let raw_notary_signatures = notary_signatures
+			.values()
+			.flat_map(Base64::<Standard>::parse);
+
+		for signature in raw_notary_signatures {
+			let signature_bytes = signature.as_bytes();
+			for notary_key in raw_notary_keys.iter().copied() {
+				let verified = ruma::signatures::verify_canonical_json_bytes(
+					&SigningKeyAlgorithm::Ed25519,
+					notary_key,
+					signature_bytes,
+					for_verify.as_bytes(),
+				);
+				if verified.is_ok() {
+					trace!(
+						"Valid signature from {notary_name} present on signing keys response \
+						 for {}",
+						server_keys.server_name
+					);
+					return Ok(());
+				}
+			}
+		}
+
+		Err!(Request(Forbidden(
+			"No valid signature from {notary_name} present on signing keys response"
+		)))
+	}
+
+	fn strip_extraneous_signatures(
+		canonical: &mut CanonicalJsonObject,
+		origin: &ServerName,
+		notaries: &[&ServerName],
+	) {
+		canonical.entry("signatures".to_owned()).and_modify(|sigs| {
+			sigs.as_object_mut().map(|s| {
+				s.retain(|server_name, _| {
+					let Ok(server_name) = ServerName::parse(server_name) else { return false };
+					origin == server_name || notaries.iter().any(|ns| *ns == server_name)
+				});
+				Some(s)
+			});
+		});
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use ruma::{
+		ServerSigningKeyId,
+		api::federation::discovery::get_server_keys::v2::Response as DirectKeysResponse,
+		owned_server_name, serde::Raw, server_name,
+	};
+
+	use super::*;
+	use crate::server_keys::Service;
+
+	#[test]
+	fn verify_direct_server_keys_response() {
+		let response_data = serde_json::from_value(serde_json::json!(
+			{
+			  "old_verify_keys": {},
+			  "server_name": "continuwuity.org",
+			  "signatures": {
+				"continuwuity.org": {
+				  "ed25519:PwHlNsFu": "YUA+lUu5tQcRhd44bxFcY+WRu8X9nLXT8FUpY44O48/oKF/9rWUhqJGw5uZeOuo2oBMAxtUmO6HT2kBI9WftBA"
+				}
+			  },
+			  "valid_until_ts": 1787690539756_i64,
+			  "verify_keys": {
+				"ed25519:PwHlNsFu": {
+				  "key": "8eNx2s0zWW+heKAmOH5zKv/nCPkEpraDJfGHxDu6hFI"
+				}
+			  }
+			}
+		)).map(Raw::from_json).unwrap();
+
+		Service::verify_server_keys_response(&response_data.deserialize().unwrap(), None)
+			.expect("Self-signature of real key response is valid");
+	}
+
+	#[test]
+	fn verify_direct_server_keys_response_ignores_extraneous_signatures() {
+		let response_data = serde_json::from_value(serde_json::json!(
+			{
+			  "old_verify_keys": {},
+			  "server_name": "continuwuity.org",
+			  "signatures": {
+				"continuwuity.org": {
+				  "ed25519:PwHlNsFu": "YUA+lUu5tQcRhd44bxFcY+WRu8X9nLXT8FUpY44O48/oKF/9rWUhqJGw5uZeOuo2oBMAxtUmO6HT2kBI9WftBA"
+				},
+				"example.org": {
+				  "ed25519:1": "yadda yadda"
+				}
+			  },
+			  "valid_until_ts": 1787690539756_i64,
+			  "verify_keys": {
+				"ed25519:PwHlNsFu": {
+				  "key": "8eNx2s0zWW+heKAmOH5zKv/nCPkEpraDJfGHxDu6hFI"
+				}
+			  }
+			}
+		)).map(Raw::from_json).unwrap();
+
+		Service::verify_server_keys_response(&response_data.deserialize().unwrap(), None)
+			.expect("Self-signature of real key response is valid");
+	}
+
+	#[test]
+	fn verify_notary_response_checks_valid_notary_signature() {
+		let response_data = serde_json::from_value(serde_json::json!(
+			{
+			  "old_verify_keys": {},
+			  "server_name": "continuwuity.org",
+			  "signatures": {
+				"continuwuity.org": {
+				  "ed25519:PwHlNsFu": "4YjHGqJ0HN93tvev4kx0qx/3w/h1nlEkFZni15h+5MZ91xUhQNa0BzqdraPU4/KJoXVcOQpmR+hsCVYb2/WlAQ"
+				},
+				"starstruck.systems": {
+				  "ed25519:a_Herg": "G4uXHfn9lizmixPFCBRXVnyf1HB1yyrxPVbtSOdViNHRFrFF66xmEpzDPs341L/Pjxtwfndi4Hw7zIsyTRqoCQ"
+				}
+			  },
+			  "valid_until_ts": 1787691347175_i64,
+			  "verify_keys": {
+				"ed25519:PwHlNsFu": {
+				  "key": "8eNx2s0zWW+heKAmOH5zKv/nCPkEpraDJfGHxDu6hFI"
+				}
+			  }
+			}
+		)).map(Raw::from_json).unwrap();
+		let starstruck_pubkey =
+			Base64::<Standard>::parse("Z4oHuEjdUbtC5W/z/gzc5B8raxNk+8gGwbNwgxHqW8Y").unwrap();
+		Service::verify_server_keys_response(
+			&response_data.deserialize().unwrap(),
+			Some((server_name!("starstruck.systems"), &[starstruck_pubkey])),
+		)
+		.expect("Server keys from notary are valid");
+	}
+
+	#[test]
+	fn ensure_verify_server_keys_response_rejects_invalid_self_signature() {
+		// even something as small as a mismatch in the timestamp should be
+		// enough to trigger this.
+		let response_data = serde_json::from_value(serde_json::json!(
+			{
+			  "old_verify_keys": {},
+			  "server_name": "continuwuity.org",
+			  "signatures": {
+				"continuwuity.org": {
+				  "ed25519:PwHlNsFu": "YUA+lUu5tQcRhd44bxFcY+WRu8X9nLXT8FUpY44O48/oKF/9rWUhqJGw5uZeOuo2oBMAxtUmO6HT2kBI9WftBA"
+				}
+			  },
+			  "valid_until_ts": 1787690539757_i64,
+			  "verify_keys": {
+				"ed25519:PwHlNsFu": {
+				  "key": "8eNx2s0zWW+heKAmOH5zKv/nCPkEpraDJfGHxDu6hFI"
+				}
+			  }
+			}
+		)).map(Raw::from_json).unwrap();
+		Service::verify_server_keys_response(&response_data.deserialize().unwrap(), None)
+			.expect_err("Tampered response should fail self-signature check");
+	}
+
+	#[test]
+	fn ensure_mismatched_notary_signature_fails() {
+		// The signature on a notary response is typically useless. However,
+		// when additional verification is configured, it can be used to prevent
+		// an attack where the notary's TLS is intercepted, and the attacker
+		// returns a valid signing key response that is signed by the attacker
+		// with the attacker controlled key. In this case, the self-signature
+		// would pass, and the attacker would be allowed to inject an
+		// adversarial signing key into our database, allowing them to
+		// completely forge any event they ever want to as the target.
+		// However, by basically doing key pinning, this attack can be mitigated
+		// by raising the attack difficulty to requiring both TLS interception,
+		// and the target notary's private key.
+
+		// In this case, we'll use a valid self-signature, BUT incorrect
+		// notary signature.
+		// Without the additional notary validation, this would permit as TOFU.
+		let response_data = serde_json::from_value(serde_json::json!(
+			{
+			  "old_verify_keys": {},
+			  "server_name": "continuwuity.org",
+			  "signatures": {
+				"continuwuity.org": {
+				  "ed25519:PwHlNsFu": "Ov3RCRYRUEgaMPfHP4YPUvjoBKE0yVTERNlgeHpxvh/+fknGeRnIZ4e6b0n1aI+gPU04Eb1DE4gaeTFigCudAg"
+				}
+			  },
+			  "valid_until_ts": 1787690539756_i64,
+			  "verify_keys": {
+				"ed25519:PwHlNsFu": {
+				  "key": "8eNx2s0zWW+heKAmOH5zKv/nCPkEpraDJfGHxDu6hFI"
+				}
+			  }
+			}
+		)).map(Raw::from_json).unwrap();
+		Service::verify_server_keys_response(&response_data.deserialize().unwrap(), None)
+			.expect_err("Intercepted notary response should fail strict validation");
 	}
 }