fix(dashboard): replace self-calling restart logic with standalone endpoint

Akinator31 · Akinator31 · commit b16ec15c6573 · 2025-11-27T16:16:28.000+01:00
diff --git a/rustmail/src/api/handler/bot/restart.rs b/rustmail/src/api/handler/bot/restart.rs
@@ -11,46 +11,31 @@ use tokio::sync::Mutex;
 pub async fn handle_restart_bot(
     State(bot_state): State<Arc<Mutex<BotState>>>,
 ) -> (StatusCode, Json<&'static str>) {
-    let token = {
-        let state_lock = bot_state.lock().await;
-        state_lock.internal_token.clone()
-    };
-
     let state_lock = bot_state.lock().await;
 
     match state_lock.status {
         BotStatus::Stopped => {
             drop(state_lock);
 
-            let result = ping_internal("/api/bot/start", &token).await.map_err(|_| {
-                (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    Json("Failed to restart bot"),
-                )
-            });
-
-            println!("Restart result: {:?}", result);
-
-            (StatusCode::OK, Json("Bot is starting"))
+            match start_bot(bot_state).await {
+                StartBotResponse::Success(status, message) => (status, message),
+                StartBotResponse::Conflict(status, message) => (status, message),
+            }
         }
         BotStatus::Running { .. } => {
             drop(state_lock);
 
-            let _ = ping_internal("/api/bot/stop", &token).await.map_err(|_| {
-                (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    Json("Failed to stop bot"),
-                )
-            });
+            match stop_bot(bot_state.clone()).await {
+                StopBotResponse::Success(..) => {}
+                StopBotResponse::Conflict(status, message) => return (status, message),
+            }
 
             sleep(Duration::from_secs(2)).await;
 
-            let _ = ping_internal("/api/bot/start", &token).await.map_err(|_| {
-                (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    Json("Failed to stop bot"),
-                )
-            });
+            match start_bot(bot_state).await {
+                StartBotResponse::Success(..) => {}
+                StartBotResponse::Conflict(status, message) => return (status, message),
+            }
 
             (StatusCode::OK, Json("Bot is restarting"))
         }
diff --git a/rustmail/src/api/handler/bot/start.rs b/rustmail/src/api/handler/bot/start.rs
@@ -1,43 +1,16 @@
-use crate::bot::run_bot;
-use crate::prelude::config::*;
+use crate::prelude::api::*;
 use crate::prelude::types::*;
 use axum::Json;
 use axum::extract::State;
 use axum::http::StatusCode;
 use std::sync::Arc;
-use tokio::spawn;
 use tokio::sync::Mutex;
 
 pub async fn handle_start_bot(
     State(bot_state): State<Arc<Mutex<BotState>>>,
 ) -> (StatusCode, Json<&'static str>) {
-    let mut state_lock = bot_state.lock().await;
-    match state_lock.status {
-        BotStatus::Stopped => {
-            state_lock.config = load_config("config.toml");
-
-            if state_lock.config.is_none() {
-                return (StatusCode::BAD_REQUEST, Json("Missing configuration."));
-            }
-
-            let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(false);
-            let (command_tx, command_rx) = tokio::sync::mpsc::channel(32);
-            let bot_state_clone = bot_state.clone();
-
-            let handle = spawn(async move {
-                run_bot(bot_state_clone.clone(), &mut shutdown_rx, command_rx).await;
-                let mut s = bot_state_clone.lock().await;
-                s.status = BotStatus::Stopped;
-            });
-            state_lock.status = BotStatus::Running {
-                handle,
-                shutdown: shutdown_tx,
-            };
-            state_lock.command_tx = command_tx;
-
-            drop(state_lock);
-            (StatusCode::OK, Json("Bot is starting"))
-        }
-        BotStatus::Running { .. } => (StatusCode::CONFLICT, Json("Bot is already running")),
+    match start_bot(bot_state).await {
+        StartBotResponse::Success(status, message) => (status, message),
+        StartBotResponse::Conflict(status, message) => (status, message),
     }
 }
diff --git a/rustmail/src/api/handler/bot/stop.rs b/rustmail/src/api/handler/bot/stop.rs
@@ -1,4 +1,5 @@
 use crate::prelude::types::*;
+use crate::prelude::api::*;
 use axum::Json;
 use axum::extract::State;
 use axum::http::StatusCode;
@@ -8,21 +9,8 @@ use tokio::sync::Mutex;
 pub async fn handle_stop_bot(
     State(bot_state): State<Arc<Mutex<BotState>>>,
 ) -> (StatusCode, Json<&'static str>) {
-    let handle_and_shutdown = {
-        let mut state_lock = bot_state.lock().await;
-
-        match std::mem::replace(&mut state_lock.status, BotStatus::Stopped) {
-            BotStatus::Running { handle, shutdown } => Some((handle, shutdown)),
-            BotStatus::Stopped => None,
-        }
-    };
-
-    if let Some((handle, shutdown_tx)) = handle_and_shutdown {
-        let _ = shutdown_tx.send(true);
-        handle.await.unwrap();
-        (StatusCode::OK, Json("Bot stopped"))
-    } else {
-        println!("Not Starting bot stop");
-        (StatusCode::CONFLICT, Json("Bot is not running"))
+    match stop_bot(bot_state).await {
+        StopBotResponse::Success(status, message) => (status, message),
+        StopBotResponse::Conflict(status, message) => (status, message),
     }
 }
diff --git a/rustmail/src/api/middleware/auth.rs b/rustmail/src/api/middleware/auth.rs
@@ -1,17 +1,15 @@
 use crate::prelude::api::*;
 use crate::prelude::types::*;
 use axum::extract::State;
-use axum::extract::{ConnectInfo, Request};
+use axum::extract::Request;
 use axum::middleware::Next;
 use axum::response::{IntoResponse, Response};
 use axum_extra::extract::CookieJar;
 use chrono::Utc;
 use hyper::StatusCode;
 use serenity::all::{GuildId, UserId};
 use sqlx::{Row, query};
-use std::net::SocketAddr;
 use std::sync::Arc;
-use subtle::ConstantTimeEq;
 use tokio::sync::Mutex;
 
 async fn check_user_with_bot(bot_state: Arc<Mutex<BotState>>, user_id: &str) -> bool {
@@ -84,25 +82,10 @@ async fn verify_user(user_id: &str, guild_id: u64, bot_state: Arc<Mutex<BotState
 
 pub async fn auth_middleware(
     State(bot_state): State<Arc<Mutex<BotState>>>,
-    ConnectInfo(addr): ConnectInfo<SocketAddr>,
     jar: CookieJar,
     req: Request,
     next: Next,
 ) -> Response {
-    if addr.ip().is_loopback() {
-        if let Some(h) = req.headers().get("x-internal-call") {
-            if let Ok(s) = h.to_str() {
-                let state_lock = bot_state.lock().await;
-                let expected = state_lock.internal_token.as_bytes();
-
-                if expected.ct_eq(s.as_bytes()).unwrap_u8() == 1 {
-                    drop(state_lock);
-                    return next.run(req).await;
-                }
-            }
-        }
-    }
-
     let session_cookie = jar.get("session_id");
 
     if session_cookie.is_none() {
diff --git a/rustmail/src/api/utils/bot.rs b/rustmail/src/api/utils/bot.rs
@@ -0,0 +1,71 @@
+use std::sync::Arc;
+use axum::http::StatusCode;
+use axum::Json;
+use tokio::spawn;
+use tokio::sync::Mutex;
+use crate::bot::run_bot;
+use crate::config::load_config;
+use crate::types::{BotState, BotStatus};
+
+pub enum StartBotResponse {
+    Success(StatusCode, Json<&'static str>),
+    Conflict(StatusCode, Json<&'static str>),
+}
+
+pub enum StopBotResponse {
+    Success(StatusCode, Json<&'static str>),
+    Conflict(StatusCode, Json<&'static str>),
+}
+
+pub async fn start_bot(bot_state: Arc<Mutex<BotState>>) -> StartBotResponse
+{
+    let mut state_lock = bot_state.lock().await;
+    match state_lock.status {
+        BotStatus::Stopped => {
+            state_lock.config = load_config("config.toml");
+
+            if state_lock.config.is_none() {
+                return StartBotResponse::Conflict(StatusCode::BAD_REQUEST, Json("Missing configuration."));
+            }
+
+            let (shutdown_tx, mut shutdown_rx) = tokio::sync::watch::channel(false);
+            let (command_tx, command_rx) = tokio::sync::mpsc::channel(32);
+            let bot_state_clone = bot_state.clone();
+
+            let handle = spawn(async move {
+                run_bot(bot_state_clone.clone(), &mut shutdown_rx, command_rx).await;
+                let mut s = bot_state_clone.lock().await;
+                s.status = BotStatus::Stopped;
+            });
+            state_lock.status = BotStatus::Running {
+                handle,
+                shutdown: shutdown_tx,
+            };
+            state_lock.command_tx = command_tx;
+
+            drop(state_lock);
+            StartBotResponse::Success(StatusCode::OK, Json("Bot is starting"))
+        }
+        BotStatus::Running { .. } => StartBotResponse::Conflict(StatusCode::BAD_REQUEST, Json("Bot is already running")),
+    }
+}
+
+pub async fn stop_bot(bot_state: Arc<Mutex<BotState>>) -> StopBotResponse {
+    let handle_and_shutdown = {
+        let mut state_lock = bot_state.lock().await;
+
+        match std::mem::replace(&mut state_lock.status, BotStatus::Stopped) {
+            BotStatus::Running { handle, shutdown } => Some((handle, shutdown)),
+            BotStatus::Stopped => None,
+        }
+    };
+
+    if let Some((handle, shutdown_tx)) = handle_and_shutdown {
+        let _ = shutdown_tx.send(true);
+        handle.await.unwrap();
+        StopBotResponse::Success(StatusCode::OK, Json("Bot stopped"))
+    } else {
+        println!("Not Starting bot stop");
+        StopBotResponse::Conflict(StatusCode::CONFLICT, Json("Bot is not running"))
+    }
+}
diff --git a/rustmail/src/api/utils/mod.rs b/rustmail/src/api/utils/mod.rs
@@ -1,9 +1,9 @@
 pub mod get_user_id_from_session;
 pub mod permissions_cache;
-pub mod ping_internal;
 pub mod user_permissions;
+pub mod bot;
 
 pub use get_user_id_from_session::*;
 pub use permissions_cache::*;
-pub use ping_internal::*;
 pub use user_permissions::*;
+pub use bot::*;
diff --git a/rustmail/src/api/utils/ping_internal.rs b/rustmail/src/api/utils/ping_internal.rs
diff --git a/rustmail/src/bot.rs b/rustmail/src/bot.rs
@@ -5,8 +5,6 @@ use crate::prelude::errors::*;
 use crate::prelude::handlers::*;
 use crate::prelude::panel_commands::*;
 use crate::prelude::types::*;
-use base64::Engine;
-use rand::RngCore;
 use serenity::all::{ClientBuilder, GatewayIntents, ShardManager};
 use serenity::cache::Settings as CacheSettings;
 use serenity::prelude::TypeMapKey;
@@ -24,13 +22,9 @@ impl TypeMapKey for ShardManagerKey {
 }
 
 pub async fn init_bot_state() -> Arc<Mutex<BotState>> {
-    let pool = init_database().await.expect("An error occured!");
+    let pool = init_database().await.expect("An error occurred!");
     println!("Database connected!");
 
-    let mut bytes = [0u8; 32];
-    rand::rng().fill_bytes(&mut bytes);
-    let token = base64::engine::general_purpose::STANDARD.encode(&bytes);
-
     let config = load_config("config.toml");
 
     let (command_tx, _command_rx) = tokio::sync::mpsc::channel(32);
@@ -41,7 +35,6 @@ pub async fn init_bot_state() -> Arc<Mutex<BotState>> {
         db_pool: Some(pool),
         command_tx: command_tx.clone(),
         bot_http: None,
-        internal_token: token,
     };
 
     Arc::new(Mutex::new(bot_state))
diff --git a/rustmail/src/types/bot.rs b/rustmail/src/types/bot.rs
@@ -25,5 +25,4 @@ pub struct BotState {
     pub db_pool: Option<sqlx::SqlitePool>,
     pub command_tx: tokio::sync::mpsc::Sender<BotCommand>,
     pub bot_http: Option<Arc<Http>>,
-    pub internal_token: String,
 }

Original file line number	Diff line number	Diff line change
`@@ -25,5 +25,4 @@ pub struct BotState {`
`25`	`25`	`pub db_pool: Option<sqlx::SqlitePool>,`
`26`	`26`	`pub command_tx: tokio::sync::mpsc::Sender<BotCommand>,`
`27`	`27`	`pub bot_http: Option<Arc<Http>>,`
`28`		`- pub internal_token: String,`
`29`	`28`	`}`