feat: Continuwuity notarizes the notary

src/service/server_keys/mod.rs

@@ -5,30 +5,42 @@ mod sign;
 mod util;
 mod verify;
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::{
+	collections::{BTreeMap, HashMap},
+	sync::Arc,
+	time::{Duration, Instant},
+};
 
+use assign::assign;
+use async_trait::async_trait;
 use conduwuit::{
-	Result, Server,
-	utils::{IterStream, ReadyExt, stream::TryIgnore},
+	Result, Server, SyncRwLock,
+	utils::{IterStream, ReadyExt, stream::TryIgnore, timepoint_from_now, to_canonical_object},
 };
 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},
+	CanonicalJsonObject, MilliSecondsSinceUnixEpoch, OwnedServerName, OwnedServerSigningKeyId,
+	ServerName, ServerSigningKeyId,
+	api::federation::discovery::{OldVerifyKey, ServerSigningKeys, VerifyKey},
 	room_version_rules::RoomVersionRules,
+	serde::Raw,
 	signatures::{Ed25519KeyPair, PublicKeyMap, PublicKeySet},
 };
+use serde_json::value::to_raw_value;
+pub use verify::strip_extraneous_signatures;
 
 use crate::{Dep, globals, sending, server_keys::util::required_keys};
 
+const VALID_UNTIL_TS_DURATION: Duration = Duration::from_hours(12);
+
 pub struct Service {
 	keypair: Box<Ed25519KeyPair>,
 	verify_keys: VerifyKeys,
 	services: Services,
 	db: Data,
+	last_lookup: SyncRwLock<HashMap<OwnedServerName, Instant>>,
 }
 
 struct Services {
@@ -45,6 +57,7 @@ pub type VerifyKeys = BTreeMap<OwnedServerSigningKeyId, VerifyKey>;
 pub type PubKeyMap = PublicKeyMap;
 pub type PubKeys = PublicKeySet;
 
+#[async_trait]
 impl crate::Service for Service {
 	fn build(args: crate::Args<'_>) -> Result<Arc<Self>> {
 		let (keypair, verify_keys) = keypair::init(args.db)?;
@@ -61,9 +74,12 @@ impl crate::Service for Service {
 			db: Data {
 				servernamekeyid_response: args.db["servernamekeyid_response"].clone(),
 			},
+			last_lookup: SyncRwLock::new(HashMap::new()),
 		}))
 	}
 
+	async fn clear_cache(&self) { self.last_lookup.write().clear(); }
+
 	fn name(&self) -> &str { crate::service::make_name(std::module_path!()) }
 }
 
@@ -183,6 +199,8 @@ impl Service {
 	}
 
 	/// Returns the stored server signing keys responses for the origin.
+	///
+	/// Does not imply that the stored responses are still in-date.
 	pub fn signing_keys_for<'a>(
 		&'a self,
 		origin: &'a ServerName,
@@ -193,4 +211,91 @@ impl Service {
 			.ignore_err()
 			.map(|(_, v): (Ignore, ServerSigningKeys)| v)
 	}
+
+	/// Determines if the server may contact the origin server to fetch keys
+	/// when acting as a notary server. This limits origin lookups to once per
+	/// minute, which prevents amplification attacks.
+	#[must_use]
+	pub fn notary_may_contact_origin(&self, server_name: &ServerName) -> bool {
+		if server_name == self.services.server.name {
+			// We always intercept origin requests to us with
+			// build_server_keys_response, so this is equivalent to
+			// calling `/_matrix/key/v2/server`.
+			return true;
+		}
+		self.last_lookup
+			.read()
+			.get(server_name)
+			.is_none_or(|last| last.elapsed() >= Duration::from_mins(1))
+	}
+
+	/// Generates a signing key response to serve to `/_matrix/key/v2/server`.
+	pub async fn build_server_keys_response(&self) -> Result<Raw<ServerSigningKeys>> {
+		let verify_keys = BTreeMap::from([self.active_verify_key()]);
+		let old_verify_keys = match &self.services.server.config.old_verify_keys {
+			| Some(old_verify_keys) => old_verify_keys.to_owned(),
+			| None => self.discover_old_signing_keys().await,
+		};
+		let valid_until_ts = timepoint_from_now(VALID_UNTIL_TS_DURATION)
+			.map(|tp| {
+				MilliSecondsSinceUnixEpoch::from_system_time(tp)
+					.expect("key validity period must be before the heat death of the universe")
+			})
+			.expect("key validity period must be before the heat death of the universe");
+		let server_keys = assign!(
+			ServerSigningKeys::new(self.services.server.name.clone(), valid_until_ts),
+			{verify_keys, old_verify_keys}
+		);
+		let mut canonical_obj = to_canonical_object(&server_keys)?;
+		self.sign_json(&mut canonical_obj)
+			.and_then(|()| to_raw_value(&canonical_obj).map_err(Into::into))
+			.map(Raw::from_json)
+	}
+
+	async fn discover_old_signing_keys(&self) -> BTreeMap<OwnedServerSigningKeyId, OldVerifyKey> {
+		let now = MilliSecondsSinceUnixEpoch::now();
+		let mut keys = self.signing_keys_for(&self.services.server.name);
+		let mut old_keys = BTreeMap::new();
+		while let Some(resp) = keys.next().await {
+			// We need to verify that this response is still trusted based on
+			// the current configuration. It's possible we got this response
+			// from a notary who we no longer trust.
+			if resp.valid_until_ts > in_one_week() || resp.valid_until_ts < now {
+				continue;
+			}
+			let trusted = self
+				.services
+				.server
+				.config
+				.trusted_servers
+				.iter()
+				.any(|notary| {
+					Self::verify_server_keys_response(
+						&resp,
+						Some((notary.server_name(), notary.verify_keys())),
+					)
+					.is_ok()
+				});
+			if !trusted {
+				continue;
+			}
+			for (old_key_id, old_key) in resp.old_verify_keys {
+				if old_key.expired_ts > resp.valid_until_ts {
+					// This key expired in the future which is probably illegal
+					continue;
+				}
+				old_keys
+					.entry(old_key_id)
+					.and_modify(|current_old_key: &mut OldVerifyKey| {
+						// Use the lowest expiry timestamp
+						if current_old_key.expired_ts < old_key.expired_ts {
+							*current_old_key = old_key.clone();
+						}
+					})
+					.or_insert(old_key);
+			}
+		}
+
+		old_keys
+	}
 }

src/service/server_keys/request.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeMap;
+use std::{collections::BTreeMap, time::Instant};
 
 use conduwuit::{Err, Result, trace, utils::millis_since_unix_epoch, warn};
 use ruma::{
@@ -30,6 +30,23 @@ impl super::Service {
 	) -> Result<ServerSigningKeys> {
 		use get_server_keys::v2::Request;
 
+		if target == self.services.server.name {
+			return self
+				.build_server_keys_response()
+				.await
+				.and_then(|r| r.deserialize().map_err(Into::into));
+		}
+
+		// N.B. The "last lookup" is written before the request is actually made
+		// to prevent concurrent notary requests from spawning... concurrent
+		// origin requests, especially if the origin is slow.
+		// This has the downside that, if the origin is temporarily unreachable
+		// (including if we're backing off from it), the notary might take an
+		// additional minute to recover compared to the rest of the server. This
+		// is deemed acceptable.
+		self.last_lookup
+			.write()
+			.insert(target.clone(), Instant::now());
 		let server_signing_key = self
 			.services
 			.sending

src/service/server_keys/verify.rs

@@ -78,7 +78,7 @@ impl super::Service {
 		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, &[]);
+		strip_extraneous_signatures(&mut canonical, &server_keys.server_name, &[]);
 		ruma::signatures::verify_json(&pubkey_map, &canonical).map_err(Into::into)
 	}
 
@@ -100,13 +100,12 @@ impl super::Service {
 			)));
 		};
 		let mut canonical_object = to_canonical_object(notary_signatures)?;
-		Self::strip_extraneous_signatures(&mut canonical_object, &server_keys.server_name, &[
+		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 for_verify =
+			ruma::signatures::to_canonical_json_string_for_signing(&canonical_object)
+				.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
@@ -137,22 +136,22 @@ impl super::Service {
 			"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)
+pub 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)]