arminius_lib_db_keys/

server_keys.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0
//! The server key queries.

mod tests;

use std::time::SystemTime;

use arminius_lib_db::{Database, DatabaseError, ResultExt as _, TransactedDb};
use ruma::{
	KeyId, OwnedServerName, OwnedServerSigningKeyId, ServerName, ServerSigningKeyId,
	ServerSigningKeyVersion, SigningKeyAlgorithm,
	api::federation::discovery::{OldVerifyKey, ServerSigningKeys},
	serde::{Base64, Raw},
	signatures::Ed25519KeyPair,
};

/// Error related to managing server keys
#[derive(thiserror::Error, Debug)]
pub enum ServerKeyError {
	/// A database error occurred
	#[error("Database error: {0}")]
	DatabaseError(#[from] DatabaseError),
	/// An error related to the key format occurred
	#[error("Key error: {0}")]
	KeyError(#[from] ruma::signatures::Ed25519KeyPairParseError),
	/// No key was found in the database
	#[error("Key not found")]
	NotFound,
	/// A configuration error occurred
	#[error("Configuration error: {0}")]
	ConfigurationError(String),
}

impl From<sqlx::Error> for ServerKeyError {
	fn from(err: sqlx::Error) -> Self {
		ServerKeyError::DatabaseError(DatabaseError::from(err))
	}
}

/// A server key from the database, and its expiration
#[derive(Debug, Eq, PartialEq)]
pub struct ServerKey {
	/// The actual key
	pub key: Vec<u8>,
	/// Time at which the key will/did expire
	///
	/// With federation, you often need to check the expiration of the key.
	/// Incoming federation requests should not be accepted with expired keys.
	/// But old messages may be signed with keys that have since been expired,
	/// and that's completely fine.
	pub expires_at: chrono::DateTime<chrono::Utc>,
}

/// Queries for managing server signing keys
// Queries are usually used as extension traits so suggest `import as _`.
#[rust_analyzer::prefer_underscore_import]
pub trait ServerKeyQueries: Database {
	/// Fetches the server key from the database, or an error if it does not
	/// exist
	#[tracing::instrument(skip(self))]
	async fn fetch_active_server_key(&mut self) -> Result<Ed25519KeyPair, ServerKeyError> {
		let res = sqlx::query!(
			"
		SELECT
			key_pair, key_version
		FROM
			server_private_keys
		WHERE
			key_pair IS NOT NULL
		"
		)
		.fetch_optional(self.connection().await?)
		.await?
		.ok_or(ServerKeyError::NotFound)?;

		Ed25519KeyPair::from_der(
			&res.key_pair.ok_or(ServerKeyError::ConfigurationError(
				"Signing key in database invalid".to_owned(),
			))?[..],
			res.key_version,
		)
		.map_err(ServerKeyError::from)
	}

	/// Gets all previous signing keys
	#[tracing::instrument(skip(self))]
	async fn fetch_expired_verify_keys(
		&mut self,
	) -> Result<Vec<(OwnedServerSigningKeyId, OldVerifyKey)>, ServerKeyError> {
		Ok(sqlx::query!(
			"SELECT expired_at, verify_key, key_version FROM server_private_keys WHERE key_pair IS NULL;"
		)
		.fetch_all(self.connection().await?)
		.await?
		.iter()
		.map(|row| {
			let expired_at = row.expired_at.ok_or(ServerKeyError::ConfigurationError(
				"Expired key does not have expiry timestamp set".to_owned(),
			))?;

			let date =
				ruma::MilliSecondsSinceUnixEpoch::from_system_time(SystemTime::from(expired_at))
					.ok_or(ServerKeyError::ConfigurationError(
						"Invalid timestamp for expired key".to_owned(),
					))?;
			let b64 = Base64::new(row.verify_key.clone());
			let name = ServerSigningKeyId::from_parts(
				SigningKeyAlgorithm::Ed25519,
				&ServerSigningKeyVersion::parse(row.key_version.clone()).map_err(|_| {
					ServerKeyError::ConfigurationError("Invalid server key version".to_owned())
				})?,
			);
			let verify_key = OldVerifyKey::new(date, b64);
			Ok(Some((name, verify_key)))
		})
		.collect::<Result<Vec<_>, ServerKeyError>>()?
		.into_iter()
		.flatten()
		.collect::<Vec<_>>())
	}

	/// Fetch foreign server keys
	#[tracing::instrument(skip(self))]
	async fn fetch_foreign_server_keys(
		&mut self,
		server_names: &ServerName,
		key_id: &KeyId<SigningKeyAlgorithm, ServerSigningKeyVersion>,
	) -> Result<Option<ServerKey>, DatabaseError> {
		let algorithm = key_id.algorithm();
		Ok(sqlx::query!(
			"
			SELECT
			    public_key,
			    expires_at
			FROM
			    server_public_keys
			WHERE
			    server_name = $1
			    AND key_algorithm = $2
			    AND key_name = $3
			",
			server_names.as_str(),
			algorithm.as_str(),
			key_id.key_name().as_str(),
		)
		.fetch_optional(self.connection().await?)
		.await?
		.map(|row| ServerKey { key: row.public_key, expires_at: row.expires_at }))
	}

	/// Upsert a server key
	///
	/// If the key already exists, update the `expires_at` value
	/// If another key exists under the same HS, but with the same name, returns
	/// an error. This is a serious misconfiguration from the other HS, and not
	/// really something we can handle in a sane manner.
	/// The spec does not give any instructions for handling this either.
	#[tracing::instrument(skip(self))]
	async fn upsert_server_key(
		&mut self,
		server_name: &ServerName,
		key_id: &KeyId<SigningKeyAlgorithm, ServerSigningKeyVersion>,
		server_key: &ServerKey,
	) -> Result<(), DatabaseError> {
		let algorithm = key_id.algorithm();
		let result = sqlx::query!(
			"
			INSERT INTO server_public_keys (server_name, key_algorithm, key_name, public_key, expires_at)
			    VALUES ($1, $2, $3, $4, $5)
			ON CONFLICT (server_name, key_name, key_algorithm)
			    DO UPDATE SET
			        expires_at = excluded.expires_at
			    WHERE
			        server_public_keys.public_key = excluded.public_key
			",
			server_name.as_str(),
			algorithm.as_str(),
			key_id.key_name().as_str(),
			server_key.key,
			server_key.expires_at,
		)
		.execute(self.connection().await?)
		.await?;
		if result.rows_affected() == 0 {
			Err(DatabaseError::InvalidData(format!(
				"A key {key_id} for the server {server_name} already exists in the database, but now another key with the same ID is being written",
			)))
		} else {
			Ok(())
		}
	}

	/// Fetch a cached key response
	#[tracing::instrument(skip(self))]
	async fn fetch_key_response(
		&mut self,
		server_names: &[OwnedServerName],
	) -> Result<Vec<Raw<ServerSigningKeys>>, DatabaseError> {
		let names: Vec<String> = server_names.iter().map(ToString::to_string).collect();
		sqlx::query!(
			"
			SELECT
			    response
			FROM
			    server_key_responses
			WHERE
			    server_name IN (
			        SELECT
			            *
			        FROM
			            UNNEST($1::text[]))
			",
			&names[..]
		)
		.fetch_all(self.connection().await?)
		.await?
		.into_iter()
		.map(|row| {
			Raw::<ServerSigningKeys>::from_json_string(row.response)
				.invalid_data("Could not create raw value")
		})
		.collect::<Result<Vec<_>, _>>()
	}

	/// Idempotently inserts a signing key coming from a homeserver import.
	///
	/// `key_id` is the full key id (e.g. `"ed25519:1"`); `key_pair_der` is
	/// the PKCS#8 v1 DER encoding of the key pair; `verify_key` is the 32-byte
	/// raw ed25519 public key.
	///
	/// The `private_key` field in the export is the unpadded standard-base64
	/// encoding of the raw 32-byte ed25519 seed, matching Synapse's signing-key
	/// file format. The importer decodes that seed, wraps it in a PKCS#8 v1
	/// header to produce `key_pair_der`, and derives `verify_key` from it
	/// before calling this method.
	///
	/// Returns `true` if a new row was inserted, `false` if the key was already
	/// present (idempotent; re-importing the same key is a no-op).
	#[tracing::instrument(skip(self, key_pair_der, verify_key))]
	async fn import_private_key(
		&mut self,
		key_id: &KeyId<SigningKeyAlgorithm, ServerSigningKeyVersion>,
		key_pair_der: &[u8],
		verify_key: &[u8],
	) -> Result<bool, DatabaseError> {
		let result = sqlx::query!(
			r#"
			INSERT INTO server_private_keys (key_version, key_pair, verify_key)
				VALUES ($1, $2, $3)
			ON CONFLICT (key_version)
				DO NOTHING
			"#,
			key_id.as_str(),
			key_pair_der,
			verify_key,
		)
		.execute(self.connection().await?)
		.await?;
		Ok(result.rows_affected() > 0)
	}

	/// Upsert a server key response from another HS
	///
	/// Replaces any previous response from the same HS
	#[tracing::instrument(skip(self))]
	async fn upsert_key_response(
		&mut self,
		server_name: &ServerName,
		key_response: Raw<ServerSigningKeys>,
	) -> Result<(), DatabaseError> {
		// Because a new request is only made when `expires_at` is in the past,
		// every new request made to the server key endpoint will have a
		// different value of `expires_at`, which needs to be in the future,
		// giving us a new response, which means that
		// the value of `response` will change each time this query is called.
		sqlx::query!(
			"
			INSERT INTO server_key_responses (server_name, response)
			    VALUES ($1, $2)
			ON CONFLICT (server_name)
			    DO UPDATE SET
			        response = excluded.response
			",
			server_name.as_str(),
			key_response.json().get()
		)
		.execute(self.connection().await?)
		.await?;
		Ok(())
	}
}

/// Server keys related queries that must operate within a transaction
pub trait TransactedServerKeyQueries: ServerKeyQueries {
	/// Deletes all signing keys from the database, and generates a new one
	async fn rotate_server_keys(&mut self) -> Result<(), ServerKeyError> {
		// Question: Should we generally be using zeroed around our keys?

		// Generating a keypair in ruma necessarily consists of creating a DER
		// certificate, and then deserializing it

		let der = Ed25519KeyPair::generate();

		// We deserialize the recently created keypair, only in order to get the
		// first eight bytes of the public key, in order to use it as a name.
		let tmp_key_pair = Ed25519KeyPair::from_der(der.as_ref(), String::new()).map_err(|_| {
			ServerKeyError::ConfigurationError("Server keys have invalid format".to_owned())
		})?;
		let name: String = bs58::encode(tmp_key_pair.public_key())
			.with_alphabet(bs58::Alphabet::BITCOIN)
			.into_string()
			.chars()
			.take(8)
			.collect();

		sqlx::query!(
			"UPDATE server_private_keys SET key_pair = NULL, expired_at = NOW() WHERE key_pair IS NOT NULL;"
		)
		.execute(&mut *self.connection().await?)
		.await?;

		let der_ref: &[u8] = der.as_ref();

		sqlx::query!(
			"INSERT INTO server_private_keys (key_version, key_pair, verify_key) VALUES ($1, $2, $3)",
			name,
			der_ref,
			&tmp_key_pair.public_key(),
		)
		.execute(&mut *self.connection().await?)
		.await?;
		Ok(())
	}

	/// If there is no signing key in the database, generate one and put it
	/// there.
	async fn generate_key_if_not_exists(&mut self) -> Result<(), ServerKeyError> {
		if self.fetch_active_server_key().await.is_ok() {
			Ok(())
		} else {
			self.rotate_server_keys().await
		}
	}
}

impl<T: Database> ServerKeyQueries for T {}

impl<'c> TransactedServerKeyQueries for TransactedDb<'c> {}