From e14b7f28f2c4d5be9e40ebbfaa8ec5ac169596cb Mon Sep 17 00:00:00 2001
From: Lambda <csgl@xiretza.xyz>
Date: Fri, 20 Sep 2024 12:05:24 +0000
Subject: [PATCH] Implement federation self-test

---
 book/changelog.md |  2 ++
 src/cli/serve.rs  | 39 ++++++++++++++++++++++++++++++++++++---
 src/config.rs     |  2 ++
 src/error.rs      |  6 ++++++
 4 files changed, 46 insertions(+), 3 deletions(-)

diff --git a/book/changelog.md b/book/changelog.md
index 66506096..cd3af642 100644
--- a/book/changelog.md
+++ b/book/changelog.md
@@ -254,3 +254,5 @@ This will be the first release of Grapevine since it was forked from Conduit
 21. Sending SIGHUP to the grapevine process now reloads TLS certificates from
     disk.
     ([!97](https://gitlab.computer.surgery/matrix/grapevine-fork/-/merge_requests/97))
+22. Added a federation self-test, perfomed automatically on startup.
+    ([!106](https://gitlab.computer.surgery/matrix/grapevine/-/merge_requests/106))
diff --git a/src/cli/serve.rs b/src/cli/serve.rs
index bfde80cf..0991066b 100644
--- a/src/cli/serve.rs
+++ b/src/cli/serve.rs
@@ -22,6 +22,7 @@ use ruma::api::{
         error::{Error as RumaError, ErrorBody, ErrorKind},
         uiaa::UiaaResponse,
     },
+    federation::discovery::get_server_version,
     IncomingRequest,
 };
 use tokio::{signal, task::JoinSet};
@@ -38,7 +39,8 @@ use crate::{
     api::{
         client_server,
         ruma_wrapper::{Ar, Ra},
-        server_server, well_known,
+        server_server::{self, AllowLoopbackRequests, LogRequestError},
+        well_known,
     },
     config::{self, Config, ListenComponent, ListenTransport},
     database::KeyValueDatabase,
@@ -98,6 +100,30 @@ pub(crate) async fn run(args: ServeArgs) -> Result<(), error::ServeCommand> {
     Ok(())
 }
 
+#[tracing::instrument]
+async fn federation_self_test() -> Result<()> {
+    let response = server_server::send_request(
+        &services().globals.config.server_name,
+        get_server_version::v1::Request {},
+        LogRequestError::Yes,
+        AllowLoopbackRequests::Yes,
+    )
+    .await?;
+
+    if !response
+        .server
+        .as_ref()
+        .is_some_and(|s| s.name.as_deref() == Some(env!("CARGO_PKG_NAME")))
+    {
+        error!(?response, "unexpected server version");
+        return Err(Error::BadConfig(
+            "Got unexpected version from our own version endpoint",
+        ));
+    }
+
+    Ok(())
+}
+
 #[allow(clippy::too_many_lines)]
 async fn run_server() -> Result<(), error::Serve> {
     use error::Serve as Error;
@@ -226,10 +252,17 @@ async fn run_server() -> Result<(), error::Serve> {
         }
     }
 
-    set_application_state(ApplicationState::Ready);
-
     tokio::spawn(handle_signals(tls_config, handles));
 
+    if config.federation.enable && config.federation.self_test {
+        federation_self_test()
+            .await
+            .map_err(error::Serve::FederationSelfTestFailed)?;
+        debug!("Federation self-test completed successfully");
+    }
+
+    set_application_state(ApplicationState::Ready);
+
     while let Some(result) = servers.join_next().await {
         let (listen, result) =
             result.expect("should be able to join server task");
diff --git a/src/config.rs b/src/config.rs
index c1ebfeab..c980c26b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -328,6 +328,7 @@ pub(crate) struct ObservabilityConfig {
 #[serde(default)]
 pub(crate) struct FederationConfig {
     pub(crate) enable: bool,
+    pub(crate) self_test: bool,
     pub(crate) trusted_servers: Vec<OwnedServerName>,
     pub(crate) max_fetch_prev_events: u16,
     pub(crate) max_concurrent_requests: u16,
@@ -338,6 +339,7 @@ impl Default for FederationConfig {
     fn default() -> Self {
         Self {
             enable: true,
+            self_test: true,
             trusted_servers: vec![
                 OwnedServerName::try_from("matrix.org").unwrap()
             ],
diff --git a/src/error.rs b/src/error.rs
index 106484df..72bcd337 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -160,4 +160,10 @@ pub(crate) enum Serve {
 
     #[error("failed to run request listener on {1}")]
     Listen(#[source] std::io::Error, ListenConfig),
+
+    #[error(
+        "federation self-test failed (set `federation.self_test = false` in \
+         config to disable)"
+    )]
+    FederationSelfTestFailed(#[source] crate::Error),
 }