From 13ae036ca04b4ebd427444252ef9856b3028b7ac Mon Sep 17 00:00:00 2001
From: Andrei Vasiliu <whyte.vuhuni@gmail.com>
Date: Sun, 16 Jan 2022 13:52:23 +0200
Subject: [PATCH] Move and refactor admin commands into admin module

---
 src/database/admin.rs | 239 +++++++++++++++++++++++++++++++++++++++++-
 src/database/rooms.rs | 220 +-------------------------------------
 2 files changed, 240 insertions(+), 219 deletions(-)

diff --git a/src/database/admin.rs b/src/database/admin.rs
index 7d2301d9..518d7587 100644
--- a/src/database/admin.rs
+++ b/src/database/admin.rs
@@ -1,10 +1,17 @@
-use std::{convert::TryInto, sync::Arc};
+use std::{convert::TryFrom, convert::TryInto, sync::Arc, time::Instant};
 
-use crate::{pdu::PduBuilder, Database};
-use rocket::futures::{channel::mpsc, stream::StreamExt};
+use crate::{
+    error::{Error, Result},
+    pdu::PduBuilder,
+    server_server, Database, PduEvent,
+};
+use rocket::{
+    futures::{channel::mpsc, stream::StreamExt},
+    http::RawStr,
+};
 use ruma::{
     events::{room::message::RoomMessageEventContent, EventType},
-    UserId,
+    EventId, RoomId, RoomVersionId, UserId,
 };
 use serde_json::value::to_raw_value;
 use tokio::sync::{MutexGuard, RwLock, RwLockReadGuard};
@@ -137,3 +144,227 @@ impl Admin {
         self.sender.unbounded_send(command).unwrap();
     }
 }
+
+pub fn parse_admin_command(db: &Database, command_line: &str, body: Vec<&str>) -> AdminCommand {
+    let mut parts = command_line.split_whitespace().skip(1);
+
+    let command_name = match parts.next() {
+        Some(command) => command,
+        None => {
+            let message = "No command given. Use <code>help</code> for a list of commands.";
+            return AdminCommand::SendMessage(RoomMessageEventContent::text_html(
+                html_to_markdown(message),
+                message,
+            ));
+        }
+    };
+
+    let args: Vec<_> = parts.collect();
+
+    match try_parse_admin_command(db, command_name, args, body) {
+        Ok(admin_command) => admin_command,
+        Err(error) => {
+            let message = format!(
+                "Encountered error while handling <code>{}</code> command:\n\
+                <pre>{}</pre>",
+                command_name, error,
+            );
+
+            AdminCommand::SendMessage(RoomMessageEventContent::text_html(
+                html_to_markdown(&message),
+                message,
+            ))
+        }
+    }
+}
+
+// Helper for `RoomMessageEventContent::text_html`, which needs the content as
+// both markdown and HTML.
+fn html_to_markdown(text: &str) -> String {
+    text.replace("<p>", "")
+        .replace("</p>", "\n")
+        .replace("<pre>", "```\n")
+        .replace("</pre>", "\n```")
+        .replace("<code>", "`")
+        .replace("</code>", "`")
+        .replace("<li>", "* ")
+        .replace("</li>", "")
+        .replace("<ul>\n", "")
+        .replace("</ul>\n", "")
+}
+
+const HELP_TEXT: &'static str = r#"
+<p>The following commands are available:</p>
+<ul>
+<li><code>register_appservice</code>: Register a bridge using its registration YAML</li>
+<li><code>unregister_appservice</code>: Unregister a bridge using its ID</li>
+<li><code>list_appservices</code>: List all the currently registered bridges</li>
+<li><code>get_auth_chain</code>: Get the `auth_chain` of a PDU</li>
+<li><code>parse_pdu</code>: Parse and print a PDU from a JSON</li>
+<li><code>get_pdu</code>: Retrieve and print a PDU by ID from the Conduit database</li>
+<li><code>database_memory_usage</code>: Print database memory usage statistics</li>
+<ul>
+"#;
+
+pub fn try_parse_admin_command(
+    db: &Database,
+    command: &str,
+    args: Vec<&str>,
+    body: Vec<&str>,
+) -> Result<AdminCommand> {
+    let command = match command {
+        "help" => AdminCommand::SendMessage(RoomMessageEventContent::text_html(
+            html_to_markdown(HELP_TEXT),
+            HELP_TEXT,
+        )),
+        "register_appservice" => {
+            if body.len() > 2 && body[0].trim() == "```" && body.last().unwrap().trim() == "```" {
+                let appservice_config = body[1..body.len() - 1].join("\n");
+                let parsed_config = serde_yaml::from_str::<serde_yaml::Value>(&appservice_config);
+                match parsed_config {
+                    Ok(yaml) => AdminCommand::RegisterAppservice(yaml),
+                    Err(e) => AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                        format!("Could not parse appservice config: {}", e),
+                    )),
+                }
+            } else {
+                AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                    "Expected code block in command body.",
+                ))
+            }
+        }
+        "unregister_appservice" => {
+            if args.len() == 1 {
+                AdminCommand::UnregisterAppservice(args[0].to_owned())
+            } else {
+                AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                    "Missing appservice identifier",
+                ))
+            }
+        }
+        "list_appservices" => AdminCommand::ListAppservices,
+        "get_auth_chain" => {
+            if args.len() == 1 {
+                if let Ok(event_id) = EventId::parse_arc(args[0]) {
+                    if let Some(event) = db.rooms.get_pdu_json(&event_id)? {
+                        let room_id_str = event
+                            .get("room_id")
+                            .and_then(|val| val.as_str())
+                            .ok_or_else(|| Error::bad_database("Invalid event in database"))?;
+
+                        let room_id = <&RoomId>::try_from(room_id_str).map_err(|_| {
+                            Error::bad_database("Invalid room id field in event in database")
+                        })?;
+                        let start = Instant::now();
+                        let count =
+                            server_server::get_auth_chain(room_id, vec![event_id], db)?.count();
+                        let elapsed = start.elapsed();
+                        return Ok(AdminCommand::SendMessage(
+                            RoomMessageEventContent::text_plain(format!(
+                                "Loaded auth chain with length {} in {:?}",
+                                count, elapsed
+                            )),
+                        ));
+                    }
+                }
+            }
+
+            AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                "Usage: get_auth_chain <event-id>",
+            ))
+        }
+        "parse_pdu" => {
+            if body.len() > 2 && body[0].trim() == "```" && body.last().unwrap().trim() == "```" {
+                let string = body[1..body.len() - 1].join("\n");
+                match serde_json::from_str(&string) {
+                    Ok(value) => {
+                        let event_id = EventId::parse(format!(
+                            "${}",
+                            // Anything higher than version3 behaves the same
+                            ruma::signatures::reference_hash(&value, &RoomVersionId::V6)
+                                .expect("ruma can calculate reference hashes")
+                        ))
+                        .expect("ruma's reference hashes are valid event ids");
+
+                        match serde_json::from_value::<PduEvent>(
+                            serde_json::to_value(value).expect("value is json"),
+                        ) {
+                            Ok(pdu) => {
+                                AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                                    format!("EventId: {:?}\n{:#?}", event_id, pdu),
+                                ))
+                            }
+                            Err(e) => AdminCommand::SendMessage(
+                                RoomMessageEventContent::text_plain(format!(
+                                    "EventId: {:?}\nCould not parse event: {}",
+                                    event_id, e
+                                )),
+                            ),
+                        }
+                    }
+                    Err(e) => AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                        format!("Invalid json in command body: {}", e),
+                    )),
+                }
+            } else {
+                AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                    "Expected code block in command body.",
+                ))
+            }
+        }
+        "get_pdu" => {
+            if args.len() == 1 {
+                if let Ok(event_id) = EventId::parse(args[0]) {
+                    let mut outlier = false;
+                    let mut pdu_json = db.rooms.get_non_outlier_pdu_json(&event_id)?;
+                    if pdu_json.is_none() {
+                        outlier = true;
+                        pdu_json = db.rooms.get_pdu_json(&event_id)?;
+                    }
+                    match pdu_json {
+                        Some(json) => {
+                            let json_text = serde_json::to_string_pretty(&json)
+                                .expect("canonical json is valid json");
+                            AdminCommand::SendMessage(
+                                RoomMessageEventContent::text_html(
+                                    format!("{}\n```json\n{}\n```",
+                                    if outlier {
+                                        "PDU is outlier"
+                                    } else { "PDU was accepted"}, json_text),
+                                    format!("<p>{}</p>\n<pre><code class=\"language-json\">{}\n</code></pre>\n",
+                                    if outlier {
+                                        "PDU is outlier"
+                                    } else { "PDU was accepted"}, RawStr::new(&json_text).html_escape())
+                                ),
+                            )
+                        }
+                        None => AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                            "PDU not found.",
+                        )),
+                    }
+                } else {
+                    AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                        "Event ID could not be parsed.",
+                    ))
+                }
+            } else {
+                AdminCommand::SendMessage(RoomMessageEventContent::text_plain(
+                    "Usage: get_pdu <eventid>",
+                ))
+            }
+        }
+        "database_memory_usage" => AdminCommand::ShowMemoryUsage,
+        _ => {
+            let message = format!(
+                "Unrecognized command <code>{}</code>, try <code>help</code> for a list of commands.",
+                command,
+            );
+            AdminCommand::SendMessage(RoomMessageEventContent::text_html(
+                html_to_markdown(&message),
+                message,
+            ))
+        }
+    };
+
+    Ok(command)
+}
diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index 0ba6c9ba..14df8f50 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -3,13 +3,13 @@ mod edus;
 pub use edus::RoomEdus;
 
 use crate::{
+    database::admin::parse_admin_command,
     pdu::{EventHash, PduBuilder},
-    server_server, utils, Database, Error, PduEvent, Result,
+    utils, Database, Error, PduEvent, Result,
 };
 use lru_cache::LruCache;
 use regex::Regex;
 use ring::digest;
-use rocket::http::RawStr;
 use ruma::{
     api::{client::error::ErrorKind, federation},
     events::{
@@ -19,7 +19,6 @@ use ruma::{
         room::{
             create::RoomCreateEventContent,
             member::{MembershipState, RoomMemberEventContent},
-            message::RoomMessageEventContent,
             power_levels::RoomPowerLevelsEventContent,
         },
         tag::TagEvent,
@@ -40,12 +39,11 @@ use std::{
     iter,
     mem::size_of,
     sync::{Arc, Mutex, RwLock},
-    time::Instant,
 };
 use tokio::sync::MutexGuard;
 use tracing::{error, warn};
 
-use super::{abstraction::Tree, admin::AdminCommand, pusher};
+use super::{abstraction::Tree, pusher};
 
 /// The unique identifier of each state group.
 ///
@@ -1496,216 +1494,8 @@ impl Rooms {
                         let command_line = lines.next().expect("each string has at least one line");
                         let body: Vec<_> = lines.collect();
 
-                        let mut parts = command_line.split_whitespace().skip(1);
-                        if let Some(command) = parts.next() {
-                            let args: Vec<_> = parts.collect();
-
-                            match command {
-                                "register_appservice" => {
-                                    if body.len() > 2
-                                        && body[0].trim() == "```"
-                                        && body.last().unwrap().trim() == "```"
-                                    {
-                                        let appservice_config = body[1..body.len() - 1].join("\n");
-                                        let parsed_config = serde_yaml::from_str::<serde_yaml::Value>(
-                                            &appservice_config,
-                                        );
-                                        match parsed_config {
-                                            Ok(yaml) => {
-                                                db.admin
-                                                    .send(AdminCommand::RegisterAppservice(yaml));
-                                            }
-                                            Err(e) => {
-                                                db.admin.send(AdminCommand::SendMessage(
-                                                    RoomMessageEventContent::text_plain(format!(
-                                                        "Could not parse appservice config: {}",
-                                                        e
-                                                    )),
-                                                ));
-                                            }
-                                        }
-                                    } else {
-                                        db.admin.send(AdminCommand::SendMessage(
-                                            RoomMessageEventContent::text_plain(
-                                                "Expected code block in command body.",
-                                            ),
-                                        ));
-                                    }
-                                }
-                                "unregister_appservice" => {
-                                    if args.len() == 1 {
-                                        db.admin.send(AdminCommand::UnregisterAppservice(
-                                            args[0].to_owned(),
-                                        ));
-                                    } else {
-                                        db.admin.send(AdminCommand::SendMessage(
-                                            RoomMessageEventContent::text_plain(
-                                                "Missing appservice identifier",
-                                            ),
-                                        ));
-                                    }
-                                }
-                                "list_appservices" => {
-                                    db.admin.send(AdminCommand::ListAppservices);
-                                }
-                                "get_auth_chain" => {
-                                    if args.len() == 1 {
-                                        if let Ok(event_id) = EventId::parse_arc(args[0]) {
-                                            if let Some(event) = db.rooms.get_pdu_json(&event_id)? {
-                                                let room_id_str = event
-                                                    .get("room_id")
-                                                    .and_then(|val| val.as_str())
-                                                    .ok_or_else(|| {
-                                                        Error::bad_database(
-                                                            "Invalid event in database",
-                                                        )
-                                                    })?;
-
-                                                let room_id = <&RoomId>::try_from(room_id_str)
-                                                    .map_err(|_| Error::bad_database("Invalid room id field in event in database"))?;
-                                                let start = Instant::now();
-                                                let count = server_server::get_auth_chain(
-                                                    room_id,
-                                                    vec![event_id],
-                                                    db,
-                                                )?
-                                                .count();
-                                                let elapsed = start.elapsed();
-                                                db.admin.send(AdminCommand::SendMessage(
-                                                    RoomMessageEventContent::text_plain(format!(
-                                                        "Loaded auth chain with length {} in {:?}",
-                                                        count, elapsed
-                                                    )),
-                                                ));
-                                            }
-                                        }
-                                    }
-                                }
-                                "parse_pdu" => {
-                                    if body.len() > 2
-                                        && body[0].trim() == "```"
-                                        && body.last().unwrap().trim() == "```"
-                                    {
-                                        let string = body[1..body.len() - 1].join("\n");
-                                        match serde_json::from_str(&string) {
-                                            Ok(value) => {
-                                                let event_id = EventId::parse(format!(
-                                                    "${}",
-                                                    // Anything higher than version3 behaves the same
-                                                    ruma::signatures::reference_hash(
-                                                        &value,
-                                                        &RoomVersionId::V6
-                                                    )
-                                                    .expect("ruma can calculate reference hashes")
-                                                ))
-                                                .expect(
-                                                    "ruma's reference hashes are valid event ids",
-                                                );
-
-                                                match serde_json::from_value::<PduEvent>(
-                                                    serde_json::to_value(value)
-                                                        .expect("value is json"),
-                                                ) {
-                                                    Ok(pdu) => {
-                                                        db.admin.send(AdminCommand::SendMessage(
-                                                            RoomMessageEventContent::text_plain(
-                                                                format!(
-                                                                    "EventId: {:?}\n{:#?}",
-                                                                    event_id, pdu
-                                                                ),
-                                                            ),
-                                                        ));
-                                                    }
-                                                    Err(e) => {
-                                                        db.admin.send(AdminCommand::SendMessage(
-                                                            RoomMessageEventContent::text_plain(
-                                                                format!("EventId: {:?}\nCould not parse event: {}", event_id, e),
-                                                            ),
-                                                        ));
-                                                    }
-                                                }
-                                            }
-                                            Err(e) => {
-                                                db.admin.send(AdminCommand::SendMessage(
-                                                    RoomMessageEventContent::text_plain(format!(
-                                                        "Invalid json in command body: {}",
-                                                        e
-                                                    )),
-                                                ));
-                                            }
-                                        }
-                                    } else {
-                                        db.admin.send(AdminCommand::SendMessage(
-                                            RoomMessageEventContent::text_plain(
-                                                "Expected code block in command body.",
-                                            ),
-                                        ));
-                                    }
-                                }
-                                "get_pdu" => {
-                                    if args.len() == 1 {
-                                        if let Ok(event_id) = EventId::parse(args[0]) {
-                                            let mut outlier = false;
-                                            let mut pdu_json =
-                                                db.rooms.get_non_outlier_pdu_json(&event_id)?;
-                                            if pdu_json.is_none() {
-                                                outlier = true;
-                                                pdu_json = db.rooms.get_pdu_json(&event_id)?;
-                                            }
-                                            match pdu_json {
-                                                Some(json) => {
-                                                    let json_text =
-                                                        serde_json::to_string_pretty(&json)
-                                                            .expect("canonical json is valid json");
-                                                    db.admin.send(AdminCommand::SendMessage(
-                                                        RoomMessageEventContent::text_html(
-                                                            format!("{}\n```json\n{}\n```",
-                                                            if outlier {
-                                                                "PDU is outlier"
-                                                            } else { "PDU was accepted"}, json_text),
-                                                            format!("<p>{}</p>\n<pre><code class=\"language-json\">{}\n</code></pre>\n",
-                                                            if outlier {
-                                                                "PDU is outlier"
-                                                            } else { "PDU was accepted"}, RawStr::new(&json_text).html_escape())
-                                                        ),
-                                                    ));
-                                                }
-                                                None => {
-                                                    db.admin.send(AdminCommand::SendMessage(
-                                                        RoomMessageEventContent::text_plain(
-                                                            "PDU not found.",
-                                                        ),
-                                                    ));
-                                                }
-                                            }
-                                        } else {
-                                            db.admin.send(AdminCommand::SendMessage(
-                                                RoomMessageEventContent::text_plain(
-                                                    "Event ID could not be parsed.",
-                                                ),
-                                            ));
-                                        }
-                                    } else {
-                                        db.admin.send(AdminCommand::SendMessage(
-                                            RoomMessageEventContent::text_plain(
-                                                "Usage: get_pdu <eventid>",
-                                            ),
-                                        ));
-                                    }
-                                }
-                                "database_memory_usage" => {
-                                    db.admin.send(AdminCommand::ShowMemoryUsage);
-                                }
-                                _ => {
-                                    db.admin.send(AdminCommand::SendMessage(
-                                        RoomMessageEventContent::text_plain(format!(
-                                            "Unrecognized command: {}",
-                                            command
-                                        )),
-                                    ));
-                                }
-                            }
-                        }
+                        let command = parse_admin_command(db, command_line, body);
+                        db.admin.send(command);
                     }
                 }
             }