From 07d05fa1f9d64e507f281a026bbf7fc94597cbf3 Mon Sep 17 00:00:00 2001
From: Charles Hall <charles@computer.surgery>
Date: Thu, 10 Oct 2024 12:02:54 -0700
Subject: [PATCH] add subcmd to repair some persistent state

---
 src/cli.rs        | 30 +++++++++++++++++++
 src/cli/repair.rs | 76 +++++++++++++++++++++++++++++++++++++++++++++++
 src/error.rs      | 25 ++++++++++++++++
 3 files changed, 131 insertions(+)
 create mode 100644 src/cli/repair.rs

diff --git a/src/cli.rs b/src/cli.rs
index 99641790..1e77ca62 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -13,6 +13,7 @@ use crate::{
 };
 
 mod check_config;
+mod repair;
 mod serve;
 
 /// Command line arguments
@@ -33,6 +34,21 @@ pub(crate) enum Command {
 
     /// Check the configuration file for syntax and semantic errors.
     CheckConfig(CheckConfigArgs),
+
+    /// Repair (some) persistent state (MAKE A BACKUP FIRST!)
+    ///
+    /// This subcommand is idempotent; if it exits zero, it can safely be run
+    /// again.
+    ///
+    /// If this subcommand exits nonzero, restore from a backup. Some causes
+    /// of failure can be fixed by the user and then this subcommand can be
+    /// attempted again, others may be a bug in this subcommand.
+    ///
+    /// Currently fixes the following issues:
+    ///
+    /// * <https://gitlab.computer.surgery/matrix/grapevine/-/issues/27>
+    /// * <https://gitlab.computer.surgery/matrix/grapevine/-/issues/54>
+    Repair(RepairArgs),
 }
 
 #[derive(clap::Args)]
@@ -88,6 +104,15 @@ pub(crate) struct ServeArgs {
     pub(crate) config: ConfigArg,
 }
 
+#[derive(clap::Args)]
+pub(crate) struct RepairArgs {
+    #[clap(flatten)]
+    observability: ObservabilityArgs,
+
+    #[clap(flatten)]
+    pub(crate) config: ConfigArg,
+}
+
 impl Args {
     pub(crate) async fn run(self) -> Result<(), error::Main> {
         if let Some((format, filter)) = self.command.cli_observability_args() {
@@ -99,6 +124,7 @@ impl Args {
             Command::CheckConfig(args) => {
                 check_config::run(args.config).await?;
             }
+            Command::Repair(args) => repair::run(args).await?,
         }
         Ok(())
     }
@@ -113,6 +139,10 @@ impl Command {
                 args.observability.log_format,
                 args.observability.log_filter.clone(),
             )),
+            Command::Repair(args) => Some((
+                args.observability.log_format,
+                args.observability.log_filter.clone(),
+            )),
             Command::Serve(_) => None,
         }
     }
diff --git a/src/cli/repair.rs b/src/cli/repair.rs
new file mode 100644
index 00000000..91ac205f
--- /dev/null
+++ b/src/cli/repair.rs
@@ -0,0 +1,76 @@
+#![warn(clippy::missing_docs_in_private_items)]
+
+//! Implementation of the `repair` subcommand
+
+use std::error::Error;
+
+use tracing as t;
+
+use super::RepairArgs;
+use crate::{config, database::KeyValueDatabase, error, services, Services};
+
+/// Subcommand entrypoint
+pub(crate) async fn run(args: RepairArgs) -> Result<(), error::RepairCommand> {
+    use error::RepairCommand as Error;
+
+    t::info!("Repairing persistent state");
+
+    let config =
+        config::load(args.config.config).await.map_err(Error::Config)?;
+
+    let db = Box::leak(Box::new(
+        KeyValueDatabase::load_or_create(&config).map_err(Error::Database)?,
+    ));
+
+    Services::build(db, config, None)
+        .map_err(Error::InitializeServices)?
+        .install();
+
+    services().globals.err_if_server_name_changed()?;
+
+    db.apply_migrations().await.map_err(Error::Database)?;
+
+    repair_roomuserid_joined(db)
+        .map_err(|e| Error::Repair("roomuserid_joined", e))?;
+
+    t::info!("Done");
+
+    Ok(())
+}
+
+/// Repair the `roomuserid_joined` map
+#[t::instrument(skip(db))]
+fn repair_roomuserid_joined(
+    db: &KeyValueDatabase,
+) -> Result<(), Box<dyn Error>> {
+    t::info!(length = db.roomuserid_joined.iter().count(), "Old map length");
+
+    db.roomuserid_joined.clear()?;
+
+    for (k, _) in db.userroomid_joined.iter() {
+        let mut segments = k.split(|x| *x == 0xFF);
+
+        let Some(user_id_bytes) = segments.next() else {
+            t::warn!("No user ID bytes in key, skipping this pair");
+            continue;
+        };
+
+        let Some(room_id_bytes) = segments.next() else {
+            t::warn!("No room ID bytes in key, skipping this pair");
+            continue;
+        };
+
+        let count = segments.count();
+        if count != 0 {
+            t::warn!(count, "Extra segments in key, ignoring them");
+        }
+
+        let key = [room_id_bytes, user_id_bytes].join(&0xFF);
+
+        db.roomuserid_joined.insert(&key, &[])?;
+    }
+
+    t::info!(length = db.roomuserid_joined.iter().count(), "New map length");
+
+    Ok(())
+}
diff --git a/src/error.rs b/src/error.rs
index ce8ebb37..c9c29d01 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -48,6 +48,9 @@ pub(crate) enum Main {
 
     #[error(transparent)]
     CheckConfigCommand(#[from] CheckConfigCommand),
+
+    #[error(transparent)]
+    RepairCommand(#[from] RepairCommand),
 }
 
 /// Errors returned from the `serve` CLI subcommand.
@@ -85,6 +88,28 @@ pub(crate) enum CheckConfigCommand {
     Config(#[from] Config),
 }
 
+/// Errors returned from the `repair` CLI subcommand.
+// Missing docs are allowed here since that kind of information should be
+// encoded in the error messages themselves anyway.
+#[allow(missing_docs)]
+#[derive(Error, Debug)]
+pub(crate) enum RepairCommand {
+    #[error("failed to load configuration")]
+    Config(#[from] Config),
+
+    #[error("failed to load or create the database")]
+    Database(#[source] crate::utils::error::Error),
+
+    #[error("failed to initialize services")]
+    InitializeServices(#[source] crate::utils::error::Error),
+
+    #[error("`server_name` change check failed")]
+    ServerNameChanged(#[from] ServerNameChanged),
+
+    #[error("repair \"{0}\" failed, restoring from a backup is recommended")]
+    Repair(&'static str, #[source] Box<dyn std::error::Error>),
+}
+
 /// Error generated if `server_name` has changed or if checking this failed
 // Missing docs are allowed here since that kind of information should be
 // encoded in the error messages themselves anyway.