Refactor server_keys service/interface and related callsites

Signed-off-by: Jason Volk <jason@zemos.net>
Signed-off-by: strawberry <strawberry@puppygock.gay>

src/service/server_keys/acquire.rs

@@ -0,0 +1,175 @@
+use std::{
+	borrow::Borrow,
+	collections::{BTreeMap, BTreeSet},
+};
+
+use conduit::{debug, debug_warn, error, implement, result::FlatOk, warn};
+use futures::{stream::FuturesUnordered, StreamExt};
+use ruma::{
+	api::federation::discovery::ServerSigningKeys, serde::Raw, CanonicalJsonObject, OwnedServerName,
+	OwnedServerSigningKeyId, ServerName, ServerSigningKeyId,
+};
+use serde_json::value::RawValue as RawJsonValue;
+
+use super::key_exists;
+
+type Batch = BTreeMap<OwnedServerName, Vec<OwnedServerSigningKeyId>>;
+
+#[implement(super::Service)]
+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;
+}
+
+#[implement(super::Service)]
+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 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 missing = self.acquire_locals(batch).await;
+	let missing_keys = keys_count(&missing);
+	let missing_servers = missing.len();
+	if missing_servers == 0 {
+		return;
+	}
+
+	debug!("missing {missing_keys} keys for {missing_servers} servers locally");
+
+	let missing = self.acquire_origins(missing.into_iter()).await;
+	let missing_keys = keys_count(&missing);
+	let missing_servers = missing.len();
+	if missing_servers == 0 {
+		return;
+	}
+
+	debug_warn!("missing {missing_keys} keys for {missing_servers} servers unreachable");
+
+	let missing = self.acquire_notary(missing.into_iter()).await;
+	let missing_keys = keys_count(&missing);
+	let missing_servers = missing.len();
+	if missing_keys > 0 {
+		debug_warn!("still missing {missing_keys} keys for {missing_servers} servers from all notaries");
+		warn!("did not obtain {missing_keys} of {requested_keys} keys; some events may not be accepted");
+	}
+}
+
+#[implement(super::Service)]
+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
+}
+
+#[implement(super::Service)]
+async fn acquire_origins<I>(&self, batch: I) -> Batch
+where
+	I: Iterator<Item = (OwnedServerName, Vec<OwnedServerSigningKeyId>)> + Send,
+{
+	let mut requests: FuturesUnordered<_> = batch
+		.map(|(origin, key_ids)| self.acquire_origin(origin, key_ids))
+		.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
+}
+
+#[implement(super::Service)]
+async fn acquire_origin(
+	&self, origin: OwnedServerName, mut key_ids: Vec<OwnedServerSigningKeyId>,
+) -> (OwnedServerName, Vec<OwnedServerSigningKeyId>) {
+	if let Ok(server_keys) = self.server_request(&origin).await {
+		self.add_signing_keys(server_keys.clone()).await;
+		key_ids.retain(|key_id| !key_exists(&server_keys, key_id));
+	}
+
+	(origin, key_ids)
+}
+
+#[implement(super::Service)]
+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
+}
+
+#[implement(super::Service)]
+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.iter().flat_map(|(_, key_ids)| key_ids.iter()).count() }

src/service/server_keys/get.rs

@@ -0,0 +1,86 @@
+use std::borrow::Borrow;
+
+use conduit::{implement, Err, Result};
+use ruma::{api::federation::discovery::VerifyKey, CanonicalJsonObject, RoomVersionId, ServerName, ServerSigningKeyId};
+
+use super::{extract_key, PubKeyMap, PubKeys};
+
+#[implement(super::Service)]
+pub async fn get_event_keys(&self, object: &CanonicalJsonObject, version: &RoomVersionId) -> Result<PubKeyMap> {
+	use ruma::signatures::required_keys;
+
+	let required = match required_keys(object, version) {
+		Ok(required) => required,
+		Err(e) => return Err!(BadServerResponse("Failed to determine keys required to verify: {e}")),
+	};
+
+	let batch = required
+		.iter()
+		.map(|(s, ids)| (s.borrow(), ids.iter().map(Borrow::borrow)));
+
+	Ok(self.get_pubkeys(batch).await)
+}
+
+#[implement(super::Service)]
+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
+}
+
+#[implement(super::Service)]
+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
+}
+
+#[implement(super::Service)]
+pub async fn get_verify_key(&self, origin: &ServerName, key_id: &ServerSigningKeyId) -> Result<VerifyKey> {
+	if let Some(result) = self.verify_keys_for(origin).await.remove(key_id) {
+		return Ok(result);
+	}
+
+	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);
+		}
+	}
+
+	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 {
+				self.add_signing_keys(server_key.clone()).await;
+			}
+
+			for server_key in server_keys {
+				if let Some(result) = extract_key(server_key, key_id) {
+					return Ok(result);
+				}
+			}
+		}
+	}
+
+	Err!(BadServerResponse(debug_error!(
+		?key_id,
+		?origin,
+		"Failed to fetch federation signing-key"
+	)))
+}

src/service/server_keys/keypair.rs

@@ -0,0 +1,64 @@
+use std::sync::Arc;
+
+use conduit::{debug, debug_info, err, error, utils, utils::string_from_bytes, Result};
+use database::Database;
+use ruma::{api::federation::discovery::VerifyKey, serde::Base64, signatures::Ed25519KeyPair};
+
+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 verify_key = VerifyKey {
+		key: Base64::new(keypair.public_key().to_vec()),
+	};
+
+	let id = format!("ed25519:{}", keypair.version());
+	let verify_keys: VerifyKeys = [(id.try_into()?, verify_key)].into();
+
+	Ok((keypair, verify_keys))
+}
+
+fn load(db: &Arc<Database>) -> Result<Box<Ed25519KeyPair>> {
+	let (version, key) = db["global"]
+		.get_blocking(b"keypair")
+		.map(|ref val| {
+			// database deserializer is having trouble with this so it's manual for now
+			let mut elems = val.split(|&b| b == b'\xFF');
+			let vlen = elems.next().expect("invalid keypair entry").len();
+			let ver = string_from_bytes(&val[..vlen]).expect("invalid keypair version");
+			let der = val[vlen.saturating_add(1)..].to_vec();
+			debug!("Found existing Ed25519 keypair: {ver:?}");
+			(ver, der)
+		})
+		.or_else(|e| {
+			assert!(e.is_not_found(), "unexpected error fetching keypair");
+			create(db)
+		})?;
+
+	let key =
+		Ed25519KeyPair::from_der(&key, version).map_err(|e| err!("Failed to load ed25519 keypair from der: {e:?}"))?;
+
+	Ok(Box::new(key))
+}
+
+fn create(db: &Arc<Database>) -> Result<(String, Vec<u8>)> {
+	let keypair = Ed25519KeyPair::generate().map_err(|e| err!("Failed to generate new ed25519 keypair: {e:?}"))?;
+
+	let id = utils::rand::string(8);
+	debug_info!("Generated new Ed25519 keypair: {id:?}");
+
+	let value: (String, Vec<u8>) = (id, keypair.to_vec());
+	db["global"].raw_put(b"keypair", &value);
+
+	Ok(value)
+}
+
+#[inline]
+fn remove(db: &Arc<Database>) {
+	let global = &db["global"];
+	global.remove(b"keypair");
+}

src/service/server_keys/mod.rs

@@ -1,31 +1,30 @@
-use std::{
-	collections::{BTreeMap, HashMap, HashSet},
-	sync::Arc,
-	time::{Duration, SystemTime},
-};
+mod acquire;
+mod get;
+mod keypair;
+mod request;
+mod sign;
+mod verify;
 
-use conduit::{debug, debug_error, debug_warn, err, error, info, trace, warn, Err, Result};
-use futures::{stream::FuturesUnordered, StreamExt};
+use std::{collections::BTreeMap, sync::Arc, time::Duration};
+
+use conduit::{implement, utils::time::timepoint_from_now, Result};
+use database::{Deserialized, Json, Map};
 use ruma::{
-	api::federation::{
-		discovery::{
-			get_remote_server_keys,
-			get_remote_server_keys_batch::{self, v2::QueryCriteria},
-			get_server_keys,
-		},
-		membership::create_join_event,
-	},
-	serde::Base64,
-	CanonicalJsonObject, CanonicalJsonValue, MilliSecondsSinceUnixEpoch, OwnedServerName, OwnedServerSigningKeyId,
-	RoomVersionId, ServerName,
+	api::federation::discovery::{ServerSigningKeys, VerifyKey},
+	serde::Raw,
+	signatures::{Ed25519KeyPair, PublicKeyMap, PublicKeySet},
+	MilliSecondsSinceUnixEpoch, OwnedServerSigningKeyId, ServerName, ServerSigningKeyId,
 };
 use serde_json::value::RawValue as RawJsonValue;
-use tokio::sync::{RwLock, RwLockWriteGuard};
 
 use crate::{globals, sending, Dep};
 
 pub struct Service {
+	keypair: Box<Ed25519KeyPair>,
+	verify_keys: VerifyKeys,
+	minimum_valid: Duration,
 	services: Services,
+	db: Data,
 }
 
 struct Services {
@@ -33,546 +32,135 @@ struct Services {
 	sending: Dep<sending::Service>,
 }
 
+struct Data {
+	server_signingkeys: Arc<Map>,
+}
+
+pub type VerifyKeys = BTreeMap<OwnedServerSigningKeyId, VerifyKey>;
+pub type PubKeyMap = PublicKeyMap;
+pub type PubKeys = PublicKeySet;
+
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
+		let minimum_valid = Duration::from_secs(3600);
+		let (keypair, verify_keys) = keypair::init(args.db)?;
+
 		Ok(Arc::new(Self {
+			keypair,
+			verify_keys,
+			minimum_valid,
 			services: Services {
 				globals: args.depend::<globals::Service>("globals"),
 				sending: args.depend::<sending::Service>("sending"),
 			},
+			db: Data {
+				server_signingkeys: args.db["server_signingkeys"].clone(),
+			},
 		}))
 	}
 
 	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
 }
 
-impl Service {
-	pub async fn fetch_required_signing_keys<'a, E>(
-		&'a self, events: E, pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
-	) -> Result<()>
-	where
-		E: IntoIterator<Item = &'a BTreeMap<String, CanonicalJsonValue>> + Send,
-	{
-		let mut server_key_ids = HashMap::new();
-		for event in events {
-			for (signature_server, signature) in event
-				.get("signatures")
-				.ok_or(err!(BadServerResponse("No signatures in server response pdu.")))?
-				.as_object()
-				.ok_or(err!(BadServerResponse("Invalid signatures object in server response pdu.")))?
-			{
-				let signature_object = signature.as_object().ok_or(err!(BadServerResponse(
-					"Invalid signatures content object in server response pdu.",
-				)))?;
+#[implement(Service)]
+#[inline]
+pub fn keypair(&self) -> &Ed25519KeyPair { &self.keypair }
 
-				for signature_id in signature_object.keys() {
-					server_key_ids
-						.entry(signature_server.clone())
-						.or_insert_with(HashSet::new)
-						.insert(signature_id.clone());
-				}
-			}
-		}
+#[implement(Service)]
+async fn add_signing_keys(&self, new_keys: ServerSigningKeys) {
+	let origin = &new_keys.server_name;
 
-		if server_key_ids.is_empty() {
-			// Nothing to do, can exit early
-			trace!("server_key_ids is empty, not fetching any keys");
-			return Ok(());
-		}
+	// (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())
+		});
 
-		trace!(
-			"Fetch keys for {}",
-			server_key_ids
-				.keys()
-				.cloned()
-				.collect::<Vec<_>>()
-				.join(", ")
-		);
-
-		let mut server_keys: FuturesUnordered<_> = server_key_ids
-			.into_iter()
-			.map(|(signature_server, signature_ids)| async {
-				let fetch_res = self
-					.fetch_signing_keys_for_server(
-						signature_server.as_str().try_into().map_err(|e| {
-							(
-								signature_server.clone(),
-								err!(BadServerResponse(
-									"Invalid servername in signatures of server response pdu: {e:?}"
-								)),
-							)
-						})?,
-						signature_ids.into_iter().collect(), // HashSet to Vec
-					)
-					.await;
-
-				match fetch_res {
-					Ok(keys) => Ok((signature_server, keys)),
-					Err(e) => {
-						debug_error!(
-							"Signature verification failed: Could not fetch signing key for {signature_server}: {e}",
-						);
-						Err((signature_server, e))
-					},
-				}
-			})
-			.collect();
-
-		while let Some(fetch_res) = server_keys.next().await {
-			match fetch_res {
-				Ok((signature_server, keys)) => {
-					pub_key_map
-						.write()
-						.await
-						.insert(signature_server.clone(), keys);
-				},
-				Err((signature_server, e)) => {
-					debug_warn!("Failed to fetch keys for {signature_server}: {e:?}");
-				},
-			}
-		}
-
-		Ok(())
-	}
-
-	// Gets a list of servers for which we don't have the signing key yet. We go
-	// over the PDUs and either cache the key or add it to the list that needs to be
-	// retrieved.
-	async fn get_server_keys_from_cache(
-		&self, pdu: &RawJsonValue,
-		servers: &mut BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>,
-		_room_version: &RoomVersionId,
-		pub_key_map: &mut RwLockWriteGuard<'_, BTreeMap<String, BTreeMap<String, Base64>>>,
-	) -> Result<()> {
-		let value: CanonicalJsonObject = serde_json::from_str(pdu.get()).map_err(|e| {
-			debug_error!("Invalid PDU in server response: {pdu:#?}");
-			err!(BadServerResponse(error!("Invalid PDU in server response: {e:?}")))
-		})?;
-
-		let signatures = value
-			.get("signatures")
-			.ok_or(err!(BadServerResponse("No signatures in server response pdu.")))?
-			.as_object()
-			.ok_or(err!(BadServerResponse("Invalid signatures object in server response pdu.")))?;
-
-		for (signature_server, signature) in signatures {
-			let signature_object = signature.as_object().ok_or(err!(BadServerResponse(
-				"Invalid signatures content object in server response pdu.",
-			)))?;
-
-			let signature_ids = signature_object.keys().cloned().collect::<Vec<_>>();
-
-			let contains_all_ids =
-				|keys: &BTreeMap<String, Base64>| signature_ids.iter().all(|id| keys.contains_key(id));
-
-			let origin = <&ServerName>::try_from(signature_server.as_str()).map_err(|e| {
-				err!(BadServerResponse(
-					"Invalid servername in signatures of server response pdu: {e:?}"
-				))
-			})?;
-
-			if servers.contains_key(origin) || pub_key_map.contains_key(origin.as_str()) {
-				continue;
-			}
-
-			debug!("Loading signing keys for {origin}");
-			let result: BTreeMap<_, _> = self
-				.services
-				.globals
-				.verify_keys_for(origin)
-				.await?
-				.into_iter()
-				.map(|(k, v)| (k.to_string(), v.key))
-				.collect();
-
-			if !contains_all_ids(&result) {
-				debug_warn!("Signing key not loaded for {origin}");
-				servers.insert(origin.to_owned(), BTreeMap::new());
-			}
-
-			pub_key_map.insert(origin.to_string(), result);
-		}
-
-		Ok(())
-	}
-
-	/// Batch requests homeserver signing keys from trusted notary key servers
-	/// (`trusted_servers` config option)
-	async fn batch_request_signing_keys(
-		&self, mut servers: BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>,
-		pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
-	) -> Result<()> {
-		for server in self.services.globals.trusted_servers() {
-			debug!("Asking batch signing keys from trusted server {server}");
-			match self
-				.services
-				.sending
-				.send_federation_request(
-					server,
-					get_remote_server_keys_batch::v2::Request {
-						server_keys: servers.clone(),
-					},
-				)
-				.await
-			{
-				Ok(keys) => {
-					debug!("Got signing keys: {keys:?}");
-					let mut pkm = pub_key_map.write().await;
-					for k in keys.server_keys {
-						let k = match k.deserialize() {
-							Ok(key) => key,
-							Err(e) => {
-								warn!(
-									"Received error {e} while fetching keys from trusted server {server}: {:#?}",
-									k.into_json()
-								);
-								continue;
-							},
-						};
-
-						// TODO: Check signature from trusted server?
-						servers.remove(&k.server_name);
-
-						let result = self
-							.services
-							.globals
-							.db
-							.add_signing_key(&k.server_name, k.clone())
-							.await
-							.into_iter()
-							.map(|(k, v)| (k.to_string(), v.key))
-							.collect::<BTreeMap<_, _>>();
-
-						pkm.insert(k.server_name.to_string(), result);
-					}
-				},
-				Err(e) => error!(
-					"Failed sending batched key request to trusted key server {server} for the remote servers \
-					 {servers:?}: {e}"
-				),
-			}
-		}
-
-		Ok(())
-	}
-
-	/// Requests multiple homeserver signing keys from individual servers (not
-	/// trused notary servers)
-	async fn request_signing_keys(
-		&self, servers: BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>,
-		pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
-	) -> Result<()> {
-		debug!("Asking individual servers for signing keys: {servers:?}");
-		let mut futures: FuturesUnordered<_> = servers
-			.into_keys()
-			.map(|server| async move {
-				(
-					self.services
-						.sending
-						.send_federation_request(&server, get_server_keys::v2::Request::new())
-						.await,
-					server,
-				)
-			})
-			.collect();
-
-		while let Some(result) = futures.next().await {
-			debug!("Received new Future result");
-			if let (Ok(get_keys_response), origin) = result {
-				debug!("Result is from {origin}");
-				if let Ok(key) = get_keys_response.server_key.deserialize() {
-					let result: BTreeMap<_, _> = self
-						.services
-						.globals
-						.db
-						.add_signing_key(&origin, key)
-						.await
-						.into_iter()
-						.map(|(k, v)| (k.to_string(), v.key))
-						.collect();
-					pub_key_map.write().await.insert(origin.to_string(), result);
-				}
-			}
-			debug!("Done handling Future result");
-		}
-
-		Ok(())
-	}
-
-	pub async fn fetch_join_signing_keys(
-		&self, event: &create_join_event::v2::Response, room_version: &RoomVersionId,
-		pub_key_map: &RwLock<BTreeMap<String, BTreeMap<String, Base64>>>,
-	) -> Result<()> {
-		let mut servers: BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>> = BTreeMap::new();
-
-		{
-			let mut pkm = pub_key_map.write().await;
-
-			// Try to fetch keys, failure is okay. Servers we couldn't find in the cache
-			// will be added to `servers`
-			for pdu in event
-				.room_state
-				.state
-				.iter()
-				.chain(&event.room_state.auth_chain)
-			{
-				if let Err(error) = self
-					.get_server_keys_from_cache(pdu, &mut servers, room_version, &mut pkm)
-					.await
-				{
-					debug!(%error, "failed to get server keys from cache");
-				};
-			}
-
-			drop(pkm);
-		};
-
-		if servers.is_empty() {
-			trace!("We had all keys cached locally, not fetching any keys from remote servers");
-			return Ok(());
-		}
-
-		if self.services.globals.query_trusted_key_servers_first() {
-			info!(
-				"query_trusted_key_servers_first is set to true, querying notary trusted key servers first for \
-				 homeserver signing keys."
-			);
-
-			self.batch_request_signing_keys(servers.clone(), pub_key_map)
-				.await?;
-
-			if servers.is_empty() {
-				debug!("Trusted server supplied all signing keys, no more keys to fetch");
-				return Ok(());
-			}
-
-			debug!("Remaining servers left that the notary/trusted servers did not provide: {servers:?}");
-
-			self.request_signing_keys(servers.clone(), pub_key_map)
-				.await?;
-		} else {
-			debug!("query_trusted_key_servers_first is set to false, querying individual homeservers first");
-
-			self.request_signing_keys(servers.clone(), pub_key_map)
-				.await?;
-
-			if servers.is_empty() {
-				debug!("Individual homeservers supplied all signing keys, no more keys to fetch");
-				return Ok(());
-			}
-
-			debug!("Remaining servers left the individual homeservers did not provide: {servers:?}");
-
-			self.batch_request_signing_keys(servers.clone(), pub_key_map)
-				.await?;
-		}
-
-		debug!("Search for signing keys done");
-
-		/*if servers.is_empty() {
-			warn!("Failed to find homeserver signing keys for the remaining servers: {servers:?}");
-		}*/
-
-		Ok(())
-	}
-
-	/// Search the DB for the signing keys of the given server, if we don't have
-	/// them fetch them from the server and save to our DB.
-	#[tracing::instrument(skip_all)]
-	pub async fn fetch_signing_keys_for_server(
-		&self, origin: &ServerName, signature_ids: Vec<String>,
-	) -> Result<BTreeMap<String, Base64>> {
-		let contains_all_ids = |keys: &BTreeMap<String, Base64>| signature_ids.iter().all(|id| keys.contains_key(id));
-
-		let mut result: BTreeMap<_, _> = self
-			.services
-			.globals
-			.verify_keys_for(origin)
-			.await?
-			.into_iter()
-			.map(|(k, v)| (k.to_string(), v.key))
-			.collect();
-
-		if contains_all_ids(&result) {
-			trace!("We have all homeserver signing keys locally for {origin}, not fetching any remotely");
-			return Ok(result);
-		}
-
-		// i didnt split this out into their own functions because it's relatively small
-		if self.services.globals.query_trusted_key_servers_first() {
-			info!(
-				"query_trusted_key_servers_first is set to true, querying notary trusted servers first for {origin} \
-				 keys"
-			);
-
-			for server in self.services.globals.trusted_servers() {
-				debug!("Asking notary server {server} for {origin}'s signing key");
-				if let Some(server_keys) = self
-					.services
-					.sending
-					.send_federation_request(
-						server,
-						get_remote_server_keys::v2::Request::new(
-							origin.to_owned(),
-							MilliSecondsSinceUnixEpoch::from_system_time(
-								SystemTime::now()
-									.checked_add(Duration::from_secs(3600))
-									.expect("SystemTime too large"),
-							)
-							.expect("time is valid"),
-						),
-					)
-					.await
-					.ok()
-					.map(|resp| {
-						resp.server_keys
-							.into_iter()
-							.filter_map(|e| e.deserialize().ok())
-							.collect::<Vec<_>>()
-					}) {
-					debug!("Got signing keys: {:?}", server_keys);
-					for k in server_keys {
-						self.services
-							.globals
-							.db
-							.add_signing_key(origin, k.clone())
-							.await;
-						result.extend(
-							k.verify_keys
-								.into_iter()
-								.map(|(k, v)| (k.to_string(), v.key)),
-						);
-						result.extend(
-							k.old_verify_keys
-								.into_iter()
-								.map(|(k, v)| (k.to_string(), v.key)),
-						);
-					}
-
-					if contains_all_ids(&result) {
-						return Ok(result);
-					}
-				}
-			}
-
-			debug!("Asking {origin} for their signing keys over federation");
-			if let Some(server_key) = self
-				.services
-				.sending
-				.send_federation_request(origin, get_server_keys::v2::Request::new())
-				.await
-				.ok()
-				.and_then(|resp| resp.server_key.deserialize().ok())
-			{
-				self.services
-					.globals
-					.db
-					.add_signing_key(origin, server_key.clone())
-					.await;
-
-				result.extend(
-					server_key
-						.verify_keys
-						.into_iter()
-						.map(|(k, v)| (k.to_string(), v.key)),
-				);
-				result.extend(
-					server_key
-						.old_verify_keys
-						.into_iter()
-						.map(|(k, v)| (k.to_string(), v.key)),
-				);
-
-				if contains_all_ids(&result) {
-					return Ok(result);
-				}
-			}
-		} else {
-			info!("query_trusted_key_servers_first is set to false, querying {origin} first");
-			debug!("Asking {origin} for their signing keys over federation");
-			if let Some(server_key) = self
-				.services
-				.sending
-				.send_federation_request(origin, get_server_keys::v2::Request::new())
-				.await
-				.ok()
-				.and_then(|resp| resp.server_key.deserialize().ok())
-			{
-				self.services
-					.globals
-					.db
-					.add_signing_key(origin, server_key.clone())
-					.await;
-
-				result.extend(
-					server_key
-						.verify_keys
-						.into_iter()
-						.map(|(k, v)| (k.to_string(), v.key)),
-				);
-				result.extend(
-					server_key
-						.old_verify_keys
-						.into_iter()
-						.map(|(k, v)| (k.to_string(), v.key)),
-				);
-
-				if contains_all_ids(&result) {
-					return Ok(result);
-				}
-			}
-
-			for server in self.services.globals.trusted_servers() {
-				debug!("Asking notary server {server} for {origin}'s signing key");
-				if let Some(server_keys) = self
-					.services
-					.sending
-					.send_federation_request(
-						server,
-						get_remote_server_keys::v2::Request::new(
-							origin.to_owned(),
-							MilliSecondsSinceUnixEpoch::from_system_time(
-								SystemTime::now()
-									.checked_add(Duration::from_secs(3600))
-									.expect("SystemTime too large"),
-							)
-							.expect("time is valid"),
-						),
-					)
-					.await
-					.ok()
-					.map(|resp| {
-						resp.server_keys
-							.into_iter()
-							.filter_map(|e| e.deserialize().ok())
-							.collect::<Vec<_>>()
-					}) {
-					debug!("Got signing keys: {server_keys:?}");
-					for k in server_keys {
-						self.services
-							.globals
-							.db
-							.add_signing_key(origin, k.clone())
-							.await;
-						result.extend(
-							k.verify_keys
-								.into_iter()
-								.map(|(k, v)| (k.to_string(), v.key)),
-						);
-						result.extend(
-							k.old_verify_keys
-								.into_iter()
-								.map(|(k, v)| (k.to_string(), v.key)),
-						);
-					}
-
-					if contains_all_ids(&result) {
-						return Ok(result);
-					}
-				}
-			}
-		}
-
-		Err!(BadServerResponse(warn!("Failed to find public key for server {origin:?}")))
-	}
+	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));
+}
+
+#[implement(Service)]
+async fn verify_key_exists(&self, origin: &ServerName, key_id: &ServerSigningKeyId) -> bool {
+	type KeysMap<'a> = BTreeMap<&'a ServerSigningKeyId, &'a RawJsonValue>;
+
+	let Ok(keys) = self
+		.db
+		.server_signingkeys
+		.get(origin)
+		.await
+		.deserialized::<Raw<ServerSigningKeys>>()
+	else {
+		return false;
+	};
+
+	if let Ok(Some(verify_keys)) = keys.get_field::<KeysMap<'_>>("verify_keys") {
+		if verify_keys.contains_key(key_id) {
+			return true;
+		}
+	}
+
+	if let Ok(Some(old_verify_keys)) = keys.get_field::<KeysMap<'_>>("old_verify_keys") {
+		if old_verify_keys.contains_key(key_id) {
+			return true;
+		}
+	}
+
+	false
+}
+
+#[implement(Service)]
+pub async fn verify_keys_for(&self, origin: &ServerName) -> VerifyKeys {
+	let mut keys = self
+		.signing_keys_for(origin)
+		.await
+		.map(|keys| merge_old_keys(keys).verify_keys)
+		.unwrap_or(BTreeMap::new());
+
+	if self.services.globals.server_is_ours(origin) {
+		keys.extend(self.verify_keys.clone().into_iter());
+	}
+
+	keys
+}
+
+#[implement(Service)]
+pub async fn signing_keys_for(&self, origin: &ServerName) -> Result<ServerSigningKeys> {
+	self.db.server_signingkeys.get(origin).await.deserialized()
+}
+
+#[implement(Service)]
+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")
+}
+
+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

@@ -0,0 +1,97 @@
+use std::collections::BTreeMap;
+
+use conduit::{implement, Err, Result};
+use ruma::{
+	api::federation::discovery::{
+		get_remote_server_keys,
+		get_remote_server_keys_batch::{self, v2::QueryCriteria},
+		get_server_keys, ServerSigningKeys,
+	},
+	OwnedServerName, OwnedServerSigningKeyId, ServerName, ServerSigningKeyId,
+};
+
+#[implement(super::Service)]
+pub(super) async fn batch_notary_request<'a, S, K>(
+	&self, notary: &ServerName, batch: S,
+) -> Result<Vec<ServerSigningKeys>>
+where
+	S: Iterator<Item = (&'a ServerName, K)> + Send,
+	K: Iterator<Item = &'a ServerSigningKeyId> + Send,
+{
+	use get_remote_server_keys_batch::v2::Request;
+	type RumaBatch = BTreeMap<OwnedServerName, BTreeMap<OwnedServerSigningKeyId, QueryCriteria>>;
+
+	let criteria = QueryCriteria {
+		minimum_valid_until_ts: Some(self.minimum_valid_ts()),
+	};
+
+	let mut server_keys = RumaBatch::new();
+	for (server, key_ids) in batch {
+		let entry = server_keys.entry(server.into()).or_default();
+		for key_id in key_ids {
+			entry.insert(key_id.into(), criteria.clone());
+		}
+	}
+
+	debug_assert!(!server_keys.is_empty(), "empty batch request to notary");
+	let request = Request {
+		server_keys,
+	};
+
+	self.services
+		.sending
+		.send_federation_request(notary, request)
+		.await
+		.map(|response| response.server_keys)
+		.map(|keys| {
+			keys.into_iter()
+				.map(|key| key.deserialize())
+				.filter_map(Result::ok)
+				.collect()
+		})
+}
+
+#[implement(super::Service)]
+pub async fn notary_request(&self, notary: &ServerName, target: &ServerName) -> Result<Vec<ServerSigningKeys>> {
+	use get_remote_server_keys::v2::Request;
+
+	let request = Request {
+		server_name: target.into(),
+		minimum_valid_until_ts: self.minimum_valid_ts(),
+	};
+
+	self.services
+		.sending
+		.send_federation_request(notary, request)
+		.await
+		.map(|response| response.server_keys)
+		.map(|keys| {
+			keys.into_iter()
+				.map(|key| key.deserialize())
+				.filter_map(Result::ok)
+				.collect()
+		})
+}
+
+#[implement(super::Service)]
+pub async fn server_request(&self, target: &ServerName) -> Result<ServerSigningKeys> {
+	use get_server_keys::v2::Request;
+
+	let server_signing_key = self
+		.services
+		.sending
+		.send_federation_request(target, Request::new())
+		.await
+		.map(|response| response.server_key)
+		.and_then(|key| key.deserialize().map_err(Into::into))?;
+
+	if server_signing_key.server_name != target {
+		return Err!(BadServerResponse(debug_warn!(
+			requested = ?target,
+			response = ?server_signing_key.server_name,
+			"Server responded with bogus server_name"
+		)));
+	}
+
+	Ok(server_signing_key)
+}

src/service/server_keys/sign.rs

@@ -0,0 +1,18 @@
+use conduit::{implement, Result};
+use ruma::{CanonicalJsonObject, RoomVersionId};
+
+#[implement(super::Service)]
+pub fn sign_json(&self, object: &mut CanonicalJsonObject) -> Result {
+	use ruma::signatures::sign_json;
+
+	let server_name = self.services.globals.server_name().as_str();
+	sign_json(server_name, self.keypair(), object).map_err(Into::into)
+}
+
+#[implement(super::Service)]
+pub fn hash_and_sign_event(&self, object: &mut CanonicalJsonObject, room_version: &RoomVersionId) -> Result {
+	use ruma::signatures::hash_and_sign_event;
+
+	let server_name = self.services.globals.server_name().as_str();
+	hash_and_sign_event(server_name, self.keypair(), object, room_version).map_err(Into::into)
+}

src/service/server_keys/verify.rs

@@ -0,0 +1,33 @@
+use conduit::{implement, pdu::gen_event_id_canonical_json, Err, Result};
+use ruma::{signatures::Verified, CanonicalJsonObject, CanonicalJsonValue, OwnedEventId, RoomVersionId};
+use serde_json::value::RawValue as RawJsonValue;
+
+#[implement(super::Service)]
+pub async fn validate_and_add_event_id(
+	&self, pdu: &RawJsonValue, room_version: &RoomVersionId,
+) -> Result<(OwnedEventId, CanonicalJsonObject)> {
+	let (event_id, mut value) = gen_event_id_canonical_json(pdu, room_version)?;
+	if let Err(e) = self.verify_event(&value, Some(room_version)).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))
+}
+
+#[implement(super::Service)]
+pub async fn verify_event(
+	&self, event: &CanonicalJsonObject, room_version: Option<&RoomVersionId>,
+) -> Result<Verified> {
+	let room_version = room_version.unwrap_or(&RoomVersionId::V11);
+	let keys = self.get_event_keys(event, room_version).await?;
+	ruma::signatures::verify_event(&keys, event, room_version).map_err(Into::into)
+}
+
+#[implement(super::Service)]
+pub async fn verify_json(&self, event: &CanonicalJsonObject, room_version: Option<&RoomVersionId>) -> Result {
+	let room_version = room_version.unwrap_or(&RoomVersionId::V11);
+	let keys = self.get_event_keys(event, room_version).await?;
+	ruma::signatures::verify_json(&keys, event.clone()).map_err(Into::into)
+}