arminius_lib_component/

lib.rs

// SPDX-FileCopyrightText: 2025 BWI GmbH
//
// SPDX-License-Identifier: Apache-2.0

//! Allows defining Components, which can represent one microservice in
//! Arminius. Components can also be combined, letting us host the entire
//! application as a monolith.

pub mod error;
mod message_token;
pub mod probes;
pub mod request_source;
pub mod ruma_requests;
mod tests;
use std::{future, marker::PhantomData};

pub use arminius_lib_config::ArminiusConfig;
use arminius_lib_db::DefaultDb;
use arminius_lib_nats::{
	NatsConfiguration, NatsConfigurationError, NatsCredentials, SharedNatsClient,
};
use axum::{
	Router,
	body::Body,
	extract::{FromRef, State},
	http::{HeaderName, Method, StatusCode, Uri, header},
	middleware::Next,
};
use axum_tracing_opentelemetry::middleware::OtelAxumLayer;
use color_eyre::eyre::Context as _;
use error::{InternalServerError, MatrixApiError};
pub use message_token::{MessageTokenParseError, MessagesToken};
use ruma::api::error::ErrorKind;
use tokio::signal::unix::SignalKind;
use tokio_util::sync::CancellationToken;
use tower_http::cors::{Any, CorsLayer};

use crate::probes::probe_router;

/// The state required for running a component
#[derive(Clone)]
pub struct ServerState {
	/// A connection to JetStream
	pub nats_connection: SharedNatsClient,
	/// entry point to the DB layer. TODO: remove this when the service layer is
	/// in place
	pub database: DefaultDb,
	/// Configuration for our homeserver
	pub config: ArminiusConfig,
}

impl std::fmt::Debug for ServerState {
	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
		f.debug_struct("ServerState")
			.field("nats_connection", &self.nats_connection)
			.field("database", &"(redacted)")
			.field("config", &self.config)
			.finish()
	}
}

impl ServerState {
	/// Creates a new server state for use in tests
	#[must_use]
	#[cfg(feature = "test_utils")]
	#[expect(
		clippy::expect_used,
		reason = "TODO(#2003): This exception predates us enabling `clippy::allow_attributes_without_reason`, so we don't know why it was enabled"
	)]
	pub async fn create_for_test(
		pg_pool_options: sqlx::pool::PoolOptions<sqlx::Postgres>,
		pg_connect_options: sqlx::postgres::PgConnectOptions,
	) -> ServerState {
		use std::ops::Deref as _;

		use sqlx::ConnectOptions as _;

		let pool = pg_pool_options
			.connect_with(pg_connect_options)
			.await
			.expect("Could not connect to database for tests");

		let database = DefaultDb::new(pool.clone());
		database.run_migrations().await.expect("Failed to run migrations");

		// We set all queries to be logged with log level INFO
		// We only do this after migrations have been run, in order to pollute
		// the results less
		let connect_options =
			pool.connect_options().deref().clone().log_statements(tracing::log::LevelFilter::Info);
		pool.set_connect_options(connect_options);

		let nats_configuration = NatsConfiguration::for_test().await;
		let nats_connection = nats_configuration
			.connect(
				"tests".to_owned(),
				"0.0.0".to_owned(),
				"NATS configuration just for tests".to_owned(),
			)
			.await
			.expect("Could not connect to nats");
		nats_connection.set_up_streams().await.expect("Could not set up streams");
		nats_connection
			.set_maintenance_mode(false)
			.await
			.expect("Could not disable maintenance mode");
		let config = ArminiusConfig::for_test_using_ipv6();
		ServerState { nats_connection, database, config }
	}
}

/// Create a new server state while automatically setting the component name on
/// NATS
#[macro_export]
macro_rules! new_server_state {
	() => {
		arminius_lib_component::ServerState::new(
			env!("CARGO_BIN_NAME").to_owned(),
			env!("CARGO_PKG_VERSION").to_owned(),
			env!("CARGO_PKG_DESCRIPTION").to_owned(),
		)
	};
}

impl ServerState {
	/// Creates a new DB using the provided infrastructure
	#[must_use]
	pub fn create(
		nats_connection: SharedNatsClient,
		database: DefaultDb,
		config: ArminiusConfig,
	) -> ServerState {
		Self { nats_connection, database, config }
	}
	/// Builds a new component state by reading configuration from the
	/// environment and connecting to NATS and Postgres
	pub async fn new(
		service_name: String,
		version: String,
		description: String,
	) -> Result<Self, ComponentInitializationError> {
		let nats_credentials =
			NatsCredentials::from_env().map_err(NatsConfigurationError::Credential)?;
		let database = DefaultDb::from_env().await?;
		Ok(Self {
			nats_connection: NatsConfiguration::with_credentials(nats_credentials)?
				.connect(service_name, version, description)
				.await?,
			database,
			config: ArminiusConfig::from_env()?,
		})
	}
}

/// Signals an error during component creation / initialization
#[derive(thiserror::Error, Debug)]
pub enum ComponentInitializationError {
	/// Error connecting to DB
	#[error(transparent)]
	DatabaseError(#[from] arminius_lib_db::DatabaseError),
	/// NATS configuration error
	#[error(transparent)]
	NatsConfiguration(#[from] NatsConfigurationError),
	/// Malformed Arminius Configuration
	#[error(transparent)]
	ArminiusConfigurationError(#[from] arminius_lib_config::ArminiusConfigurationError),
}

impl FromRef<ServerState> for SharedNatsClient {
	fn from_ref(component_state: &ServerState) -> SharedNatsClient {
		component_state.nats_connection.clone()
	}
}

impl FromRef<ServerState> for DefaultDb {
	fn from_ref(component_state: &ServerState) -> DefaultDb {
		component_state.database.clone()
	}
}

impl FromRef<ServerState> for ArminiusConfig {
	fn from_ref(component_state: &ServerState) -> ArminiusConfig {
		component_state.config.clone()
	}
}

/// Error returned when the user used an unknown endpoint.
#[derive(thiserror::Error, Debug, MatrixApiError)]
#[error("Endpoint not found")]
#[matrix_api_error(status_code = StatusCode::NOT_FOUND, error_kind = ErrorKind::Unrecognized, component_crate_root = crate)]
pub struct EndpointNotFound;

/// Always returns NOT_FOUND. Called as a fallback option
async fn unknown_endpoint(method: Method, uri: Uri) -> EndpointNotFound {
	tracing::info!(%uri, %method, "Got unknown endpoint");
	EndpointNotFound
}

/// Error returned when the user used a disallowed method.
#[derive(thiserror::Error, Debug, MatrixApiError)]
#[error("Method not allowed")]
#[matrix_api_error(status_code = StatusCode::METHOD_NOT_ALLOWED, error_kind = ErrorKind::Unrecognized, component_crate_root = crate)]
pub struct MethodNotAllowed;

/// Always returns METHOD_NOT_ALLOWED. Called as a method_not_allowed_fallback
async fn method_not_allowed_endpoint(method: Method, uri: Uri) -> MethodNotAllowed {
	tracing::info!(%uri, %method, "User used disallowed method");
	MethodNotAllowed
}

/// Error type for the maintenance mode middleware layer.
#[derive(MatrixApiError, thiserror::Error, Debug)]
#[matrix_api_error(component_crate_root = crate)]
pub enum MaintenanceModeLayerError {
	/// Returned when the service is in maintenance mode.
	#[error("Service in maintenance mode")]
	#[matrix_api_error(status_code = StatusCode::SERVICE_UNAVAILABLE, error_kind = ErrorKind::Unknown)]
	MaintenanceMode,

	/// Returned when an internal error occurs.
	#[error(transparent)]
	InternalServerError(#[from] InternalServerError),
}

/// Middleware layer that returns a maintenance mode error if the service is
/// currently in maintenance mode.
async fn maintenance_mode_layer(
	state: State<ServerState>,
	req: axum::http::Request<Body>,
	next: Next,
) -> Result<axum::http::Response<Body>, MaintenanceModeLayerError> {
	let in_maintenance =
		state.nats_connection.is_in_maintenance_mode().await.map_err(|_| InternalServerError)?;

	if in_maintenance {
		Err(MaintenanceModeLayerError::MaintenanceMode)
	} else {
		Ok(next.run(req).await)
	}
}

/// Represents the "microservice", which is composed out of one or more
/// components handing specific functions.
#[derive(Debug)]
#[must_use = "Initialized server but never used"]
pub struct Server<COMPONENT: Component> {
	/// State passed to all components. This is cloned for each call.
	pub state: ServerState,
	/// The router used by this component. Might be used by tests to skip doing
	/// actual network requests.
	pub router: Router,
	/// The listener on which we will listen for stuff
	pub listener: tokio::net::TcpListener,
	/// The components of this server. Not actually stored, but this is used to
	/// invoke their initialize() and run() functions.
	components: PhantomData<COMPONENT>,
}

impl<COMPONENT: Component> Server<COMPONENT> {
	/// Initialize and construct the server.
	pub async fn initialize(state: ServerState) -> color_eyre::eyre::Result<Self> {
		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		COMPONENT::initialize(&state).await?;

		let traced_router = COMPONENT::router()
			.with_state(state.clone())
			.fallback(unknown_endpoint)
			.method_not_allowed_fallback(method_not_allowed_endpoint)
			// Add spans to requests.
			.layer(OtelAxumLayer::default());

		let router = Router::new()
			.merge(traced_router)
			.layer(axum::middleware::from_fn_with_state(state.clone(), maintenance_mode_layer))
			.layer(
				// https://spec.matrix.org/v1.14/client-server-api/#web-browser-clients
				CorsLayer::new()
					.allow_headers([
						header::AUTHORIZATION,
						header::CONTENT_TYPE,
						HeaderName::from_static("x-requested-with"),
					])
					.allow_methods([
						Method::GET,
						Method::POST,
						Method::PUT,
						Method::DELETE,
						Method::OPTIONS,
					])
					.allow_origin(Any),
			)
			// We don't want probe endpoints to be traced
			.merge(probe_router(state.clone()));

		let listener = tokio::net::TcpListener::bind(state.config.addr)
			.await
			.context("Failed binding listener")?;

		Ok(Self { state, router, listener, components: PhantomData })
	}

	/// Helper to run the axum server with the router created during
	/// initialization.
	pub async fn run_axum(
		self,
		cancellation_token: CancellationToken,
	) -> color_eyre::eyre::Result<()> {
		tracing::info!(addr = %self.state.config.addr, "Starting listener");
		axum::serve(self.listener, self.router)
			.with_graceful_shutdown(cancellation_token.cancelled_owned())
			.await
			.context("Failed serving webserver")?;
		Ok(())
	}

	/// Execute the Server/microservice. Call in main. Returns after the service
	/// exited.
	pub async fn run(self) -> color_eyre::eyre::Result<()> {
		let cancellation_token = CancellationToken::new();
		// pin the future to avoid excessive stack usage
		Box::pin(self.run_with_cancellation_token(cancellation_token)).await
	}

	/// Execute the Server/microservice, and close it when the cancellation
	/// token gets used
	pub async fn run_with_cancellation_token(
		self,
		cancellation_token: CancellationToken,
	) -> color_eyre::eyre::Result<()> {
		let tc = cancellation_token.clone();
		tokio::spawn(async move {
			await_shutdown_signal().await;
			tc.cancel();
		});

		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		COMPONENT::initialize(&self.state).await?;

		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		tokio::try_join!(
			COMPONENT::run(self.state.clone(), cancellation_token.clone()),
			self.run_axum(cancellation_token),
		)?;

		Ok(())
	}
}

/// Represents a component in the microservice
pub trait Component {
	/// Returns all the routes exposed by this component.
	fn router() -> Router<ServerState> {
		Router::new()
	}
	/// Hook to run necessary setup tasks for this component.
	#[must_use]
	fn initialize(_: &ServerState) -> impl Future<Output = color_eyre::eyre::Result<()>> + Send {
		future::ready(Ok(()))
	}

	/// Hook to run background tasks and other long running async executions for
	/// this component, e.g. subscribing to a NATS stream and handling messages.
	#[must_use]
	fn run(
		_state: ServerState,
		_cancellation_token: CancellationToken,
	) -> impl Future<Output = color_eyre::eyre::Result<()>> + Send {
		future::ready(Ok(()))
	}
}

/// A component composed of two components. Initializing this component is
/// equivalent to initializing the two components one after the other. Running
/// it is equivalent to running them as a monolith.
#[derive(Debug)]
pub struct ComposedComponent<T: Component, U: Component> {
	/// Phantom data
	_1: PhantomData<T>,
	/// Phantom data
	_2: PhantomData<U>,
}

impl<T: Component, U: Component> Component for ComposedComponent<T, U> {
	fn router() -> Router<ServerState> {
		T::router().merge(U::router())
	}

	async fn initialize(database: &ServerState) -> color_eyre::eyre::Result<()> {
		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		T::initialize(database).await?;
		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		U::initialize(database).await?;
		Ok(())
	}

	async fn run(
		state: ServerState,
		cancellation_token: CancellationToken,
	) -> color_eyre::eyre::Result<()> {
		// FIXME(#1413): Address this lint.  ast-grep-ignore: missing-context
		tokio::try_join!(
			T::run(state.clone(), cancellation_token.clone()),
			U::run(state.clone(), cancellation_token.clone())
		)?;
		Ok(())
	}
}

/// A component with no endpoints - useful for healthchecks for services that
/// don't expose their own HTTP endpoints otherwise
#[derive(Debug)]
pub struct EmptyComponent;

impl Component for EmptyComponent {}

/// For a list of components `ComponentA`, `ComponentB`, `ComponentC` constructs
/// a type `ComposedComponent<ComponentA, ComposedComponent<ComponentB,
/// ComponentC>>`.
///
/// As an optimization we match subtrees up to a certain depth and build up a
/// tree like `((c,c),(c,c))` instead of `(c, (c, (c, c)))`. This about halves
/// the recursion depth for each level of additional matching we do.
#[macro_export]
macro_rules! compose_components {
        ($c1:ident,$c2:ident) => {
		::arminius_lib_component::ComposedComponent<$c1, $c2>
	};
	($c1:ident,$c2:ident,$c3:ident,$c4:ident,$c5:ident,$c6:ident,$c7:ident,$c8:ident,$c9:ident,$($c10:ident),+) => {
		::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<$c1, $c2>, ::arminius_lib_component::ComposedComponent<$c3, $c4>>, ::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<$c5, $c6>, ::arminius_lib_component::ComposedComponent<$c7, $c8>>>, compose_components!($c9,$($c10),+)>
	};
	($c1:ident,$c2:ident,$c3:ident,$c4:ident,$c5:ident,$($c6:ident),+) => {
		::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<$c1, $c2>, ::arminius_lib_component::ComposedComponent<$c3, $c4>>, compose_components!($c5,$($c6),+)>
	};
	($c1:ident,$c2:ident,$c3:ident,$($c4:ident),+) => {
		::arminius_lib_component::ComposedComponent<::arminius_lib_component::ComposedComponent<$c1, $c2>, compose_components!($c3,$($c4),+)>
	};
	($c1:ident,$c2:ident,$($c3:ident),+) => {
		::arminius_lib_component::ComposedComponent<$c1, compose_components!($c2,$($c3),+)>
	};
}

/// Future that completes when a SIGTERM or SIGINT is received, or immediately
/// if no listener can be registered.
pub async fn await_shutdown_signal() {
	let mut sigterm_signal = match tokio::signal::unix::signal(SignalKind::terminate()) {
		Ok(signal) => signal,
		Err(error) => {
			tracing::error!(%error, "Could not create listener for sigterm");
			return;
		}
	};
	tokio::select! {
		result = tokio::signal::ctrl_c() => {
			match result {
				Ok(_) => (),
				Err(error) => {
					tracing::error!(%error, "Error registering ctrl+c listener");
				}
			}
		},
		result = sigterm_signal.recv() => {
			match result {
				Some(()) => (),
				None => {
					tracing::error!("Error registering SIGTERM listener");
				}
			}
		}
	}
}