feat: Require UUID token authorization for /shutdown endpoint

build.ps1

@@ -20,7 +20,12 @@ if (Test-Path $serverExe) {
     & $serverExe --exit 2>$null
 }
 try {
-    Invoke-RestMethod -Uri "https://127.0.0.1:3000/shutdown" -Method Post -SkipCertificateCheck -ErrorAction Stop | Out-Null
+    $tokenPath = "C:\Users\reazul.ashraf\.gemini\mcp_memory\admin.token"
+    $headers = @{}
+    if (Test-Path $tokenPath) {
+        $headers["Authorization"] = "Bearer $(Get-Content -Path $tokenPath -Raw)".Trim()
+    }
+    Invoke-RestMethod -Uri "https://127.0.0.1:3000/shutdown" -Method Post -Headers $headers -SkipCertificateCheck -ErrorAction Stop | Out-Null
 } catch {
     # The response often ends prematurely because the server abruptly kills its own process during shutdown. This is expected.
 }
@@ -63,7 +68,7 @@ function Deploy-If-Needed {
     if ($EnvName -eq 'WSL') {
         # Translate source path to WSL path
         $wslSource = wsl.exe -d Ubuntu -e wslpath -u $Source
-        wsl.exe -d Ubuntu -e bash -c "mv $Dest ${Dest}.`date +%s`.old 2>/dev/null || true; cp $wslSource $Dest && chmod +x $Dest"
+        wsl.exe -d Ubuntu -e bash -c "mv $Dest ${Dest}.\$(date +%s).old 2>/dev/null || true; cp $wslSource $Dest && chmod +x $Dest"
     } else {
         if (Test-Path $Dest) {
             $rand = Get-Random

deploy.ps1

@@ -1,6 +1,11 @@
 try {
     Write-Host "Gracefully shutting down server..."
-    Invoke-RestMethod -Method Post -Uri "http://127.0.0.1:3000/shutdown" -ErrorAction SilentlyContinue
+    $tokenPath = "C:\Users\reazul.ashraf\.gemini\mcp_memory\admin.token"
+    $headers = @{}
+    if (Test-Path $tokenPath) {
+        $headers["Authorization"] = "Bearer $(Get-Content -Path $tokenPath -Raw)".Trim()
+    }
+    Invoke-RestMethod -Method Post -Uri "http://127.0.0.1:3000/shutdown" -Headers $headers -ErrorAction SilentlyContinue | Out-Null
     Start-Sleep -Seconds 2
 } catch {
     Write-Host "Server might already be down or unreachable."

server/src/main.rs

@@ -261,14 +261,33 @@ fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>>
             .route("/gate/set", post(gate_set_handler))
             .route(
                 "/shutdown",
-                post(|| async move {
+                post(
-                    std::thread::spawn(|| {
+                    |headers: axum::http::HeaderMap, State(state): State<Arc<AppState>>| async move {
-                        tracing::info!("Received shutdown request via /shutdown endpoint. Exiting process cleanly.");
+                        let token_path = state.handler.state.base_dir.join("admin.token");
-                        std::thread::sleep(std::time::Duration::from_millis(100));
+                        let expected_token = std::fs::read_to_string(&token_path)
-                        std::process::exit(0);
+                            .unwrap_or_default()
-                    });
+                            .trim()
-                    "Shutting down..."
+                            .to_string();
-                }),
+
+                        let auth_header = headers
+                            .get(axum::http::header::AUTHORIZATION)
+                            .and_then(|h| h.to_str().ok())
+                            .unwrap_or_default();
+
+                        if expected_token.is_empty() || auth_header != format!("Bearer {}", expected_token) {
+                            return (axum::http::StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
+                        }
+
+                        std::thread::spawn(|| {
+                            tracing::info!(
+                                "Received shutdown request via /shutdown endpoint. Exiting process cleanly."
+                            );
+                            std::thread::sleep(std::time::Duration::from_millis(100));
+                            std::process::exit(0);
+                        });
+                        (axum::http::StatusCode::OK, "Shutting down...").into_response()
+                    },
+                ),
             )
             .route(
                 "/",
@@ -648,26 +667,38 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let _guard = init_logging("mcp-memory-server");
     let cli = Cli::parse();
 
+    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
+        dirs::home_dir()
+            .map(|mut h| {
+                h.push(".gemini/mcp_memory");
+                h.to_string_lossy().into_owned()
+            })
+            .unwrap_or_else(|| ".gemini/mcp_memory".into())
+    });
+    let base = PathBuf::from(base_dir);
+
     if cli.exit {
         let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
-        let _ = std::process::Command::new("curl")
+        let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
-            .arg("-k")
+        let mut cmd = std::process::Command::new("curl");
-            .arg("-X")
+        cmd.arg("-k").arg("-X").arg("POST");
-            .arg("POST")
+        if !token.is_empty() {
-            .arg(format!("https://127.0.0.1:{}/shutdown", port))
+            cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
-            .output();
+        }
+        let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
         println!("Sent shutdown request to server.");
         return Ok(());
     }
 
     if cli.restart {
         let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
-        let _ = std::process::Command::new("curl")
+        let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
-            .arg("-k")
+        let mut cmd = std::process::Command::new("curl");
-            .arg("-X")
+        cmd.arg("-k").arg("-X").arg("POST");
-            .arg("POST")
+        if !token.is_empty() {
-            .arg(format!("https://127.0.0.1:{}/shutdown", port))
+            cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim()));
-            .output();
+        }
+        let _ = cmd.arg(format!("https://127.0.0.1:{}/shutdown", port)).output();
         println!("Sent shutdown request to existing server. Waiting for it to exit...");
         std::thread::sleep(std::time::Duration::from_millis(1500));
         return Ok(());
@@ -691,18 +722,14 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         }
     }
 
-    let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
-        dirs::home_dir()
-            .map(|mut h| {
-                h.push(".gemini/mcp_memory");
-                h.to_string_lossy().into_owned()
-            })
-            .unwrap_or_else(|| ".gemini/mcp_memory".into())
-    });
-    let base = PathBuf::from(base_dir);
     fs::create_dir_all(&base).expect("Failed to create store dir");
 
+    // Generate token
+    let admin_token = uuid::Uuid::new_v4().to_string();
+    std::fs::write(base.join("admin.token"), &admin_token).expect("Failed to write admin token");
+
     let redb_path = base.join("mcp_store.redb");
+
     let db = Arc::new(redb::Database::create(&redb_path).unwrap());
 
     // Ensure table exists and migrate old JSON files