fix: Don't attempt to contact origins as a notary more than once a minute

src/api/server/key.rs

@@ -177,10 +177,13 @@ pub(crate) async fn get_remote_server_keys_route(
 	body: Ruma<get_remote_server_keys::v2::Request>,
 ) -> Result<get_remote_server_keys::v2::Response> {
 	let min_valid_ts = body.minimum_valid_until_ts;
-	if let Ok(response) = services
+	if services
 		.server_keys
-		.origin_request(body.server_name.clone(), min_valid_ts)
-		.await
+		.notary_may_contact_origin(&body.server_name)
+		&& let Ok(response) = services
+			.server_keys
+			.origin_request(body.server_name.clone(), min_valid_ts)
+			.await
 	{
 		return sign_ssk(&services, response, &body.server_name)
 			.await

src/service/server_keys/mod.rs

@@ -5,18 +5,22 @@ mod sign;
 mod util;
 mod verify;
 
-use std::{collections::BTreeMap, sync::Arc};
+use std::{
+	collections::{BTreeMap, HashMap},
+	sync::Arc,
+	time::{Duration, Instant},
+};
 
 use conduwuit::{
-	Result, Server,
+	Result, Server, SyncRwLock,
 	utils::{IterStream, ReadyExt, stream::TryIgnore},
 };
 use database::{Deserialized, Ignore, Interfix, Json, Map};
 use futures::{Stream, StreamExt};
 pub use request::in_one_week;
 use ruma::{
-	CanonicalJsonObject, MilliSecondsSinceUnixEpoch, OwnedServerSigningKeyId, ServerName,
-	ServerSigningKeyId,
+	CanonicalJsonObject, MilliSecondsSinceUnixEpoch, OwnedServerName, OwnedServerSigningKeyId,
+	ServerName, ServerSigningKeyId,
 	api::federation::discovery::{ServerSigningKeys, VerifyKey},
 	room_version_rules::RoomVersionRules,
 	signatures::{Ed25519KeyPair, PublicKeyMap, PublicKeySet},
@@ -30,6 +34,7 @@ pub struct Service {
 	verify_keys: VerifyKeys,
 	services: Services,
 	db: Data,
+	last_lookup: SyncRwLock<HashMap<OwnedServerName, Instant>>,
 }
 
 struct Services {
@@ -62,6 +67,7 @@ impl crate::Service for Service {
 			db: Data {
 				servernamekeyid_response: args.db["servernamekeyid_response"].clone(),
 			},
+			last_lookup: SyncRwLock::new(HashMap::new()),
 		}))
 	}
 
@@ -194,4 +200,15 @@ 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 {
+		self.last_lookup
+			.read()
+			.get(server_name)
+			.is_none_or(|last| last.elapsed() >= Duration::from_mins(1))
+	}
 }

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,16 @@ impl super::Service {
 	) -> Result<ServerSigningKeys> {
 		use get_server_keys::v2::Request;
 
+		// 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