From aa768b5dec1338a6dbdd1b3a9a00bc6ec9d53090 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Thu, 24 Oct 2024 05:03:31 +0000
Subject: [PATCH] distill active and old keys for federation key/server
 response

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/api/server/key.rs          | 40 +++++++++++++++++++++++++---------
 src/service/server_keys/mod.rs | 17 +++++++++++++++
 2 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/src/api/server/key.rs b/src/api/server/key.rs
index 3913ce43..5284593d 100644
--- a/src/api/server/key.rs
+++ b/src/api/server/key.rs
@@ -1,10 +1,14 @@
-use std::{collections::BTreeMap, time::Duration};
+use std::{
+	collections::BTreeMap,
+	mem::take,
+	time::{Duration, SystemTime},
+};
 
 use axum::{extract::State, response::IntoResponse, Json};
 use conduit::{utils::timepoint_from_now, Result};
 use ruma::{
 	api::{
-		federation::discovery::{get_server_keys, ServerSigningKeys},
+		federation::discovery::{get_server_keys, OldVerifyKey, ServerSigningKeys},
 		OutgoingResponse,
 	},
 	serde::Raw,
@@ -21,21 +25,32 @@ use ruma::{
 // signature for the response
 pub(crate) async fn get_server_keys_route(State(services): State<crate::State>) -> Result<impl IntoResponse> {
 	let server_name = services.globals.server_name();
-	let verify_keys = services.server_keys.verify_keys_for(server_name).await;
+	let active_key_id = services.server_keys.active_key_id();
+	let mut all_keys = services.server_keys.verify_keys_for(server_name).await;
+
+	let verify_keys = all_keys
+		.remove_entry(active_key_id)
+		.expect("active verify_key is missing");
+
+	let old_verify_keys = all_keys
+		.into_iter()
+		.map(|(id, key)| (id, OldVerifyKey::new(expires_ts(), key.key)))
+		.collect();
+
 	let server_key = ServerSigningKeys {
-		verify_keys,
+		verify_keys: [verify_keys].into(),
+		old_verify_keys,
 		server_name: server_name.to_owned(),
 		valid_until_ts: valid_until_ts(),
-		old_verify_keys: BTreeMap::new(),
 		signatures: BTreeMap::new(),
 	};
 
-	let response = get_server_keys::v2::Response {
-		server_key: Raw::new(&server_key)?,
-	}
-	.try_into_http_response::<Vec<u8>>()?;
+	let server_key = Raw::new(&server_key)?;
+	let mut response = get_server_keys::v2::Response::new(server_key)
+		.try_into_http_response::<Vec<u8>>()
+		.map(|mut response| take(response.body_mut()))
+		.and_then(|body| serde_json::from_slice(&body).map_err(Into::into))?;
 
-	let mut response = serde_json::from_slice(response.body())?;
 	services.server_keys.sign_json(&mut response)?;
 
 	Ok(Json(response))
@@ -47,6 +62,11 @@ fn valid_until_ts() -> MilliSecondsSinceUnixEpoch {
 	MilliSecondsSinceUnixEpoch::from_system_time(timepoint).expect("UInt should not overflow")
 }
 
+fn expires_ts() -> MilliSecondsSinceUnixEpoch {
+	let timepoint = SystemTime::now();
+	MilliSecondsSinceUnixEpoch::from_system_time(timepoint).expect("UInt should not overflow")
+}
+
 /// # `GET /_matrix/key/v2/server/{keyId}`
 ///
 /// Gets the public signing keys of this server.
diff --git a/src/service/server_keys/mod.rs b/src/service/server_keys/mod.rs
index dc09703c..dae45a51 100644
--- a/src/service/server_keys/mod.rs
+++ b/src/service/server_keys/mod.rs
@@ -44,7 +44,9 @@ 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)?;
+		debug_assert!(verify_keys.len() == 1, "only one active verify_key supported");
 
 		Ok(Arc::new(Self {
 			keypair,
@@ -68,6 +70,21 @@ impl crate::Service for Service {
 #[inline]
 pub fn keypair(&self) -> &Ed25519KeyPair { &self.keypair }
 
+#[implement(Service)]
+#[inline]
+pub fn active_key_id(&self) -> &ServerSigningKeyId { self.active_verify_key().0 }
+
+#[implement(Service)]
+#[inline]
+pub fn active_verify_key(&self) -> (&ServerSigningKeyId, &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))
+		.expect("missing active verify_key")
+}
+
 #[implement(Service)]
 async fn add_signing_keys(&self, new_keys: ServerSigningKeys) {
 	let origin = &new_keys.server_name;