diff --git a/build.ps1 b/build.ps1 index 186064e..3dedc77 100644 --- a/build.ps1 +++ b/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 diff --git a/deploy.ps1 b/deploy.ps1 index 0da240e..4987bb8 100644 --- a/deploy.ps1 +++ b/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." diff --git a/server/src/main.rs b/server/src/main.rs index bc54dfb..b0392fb 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -261,14 +261,33 @@ fn run_server(state: Arc) -> Result<(), Box> .route("/gate/set", post(gate_set_handler)) .route( "/shutdown", - post(|| async move { - 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); - }); - "Shutting down..." - }), + post( + |headers: axum::http::HeaderMap, State(state): State>| async move { + let token_path = state.handler.state.base_dir.join("admin.token"); + let expected_token = std::fs::read_to_string(&token_path) + .unwrap_or_default() + .trim() + .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> { 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") - .arg("-k") - .arg("-X") - .arg("POST") - .arg(format!("https://127.0.0.1:{}/shutdown", port)) - .output(); + let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default(); + let mut cmd = std::process::Command::new("curl"); + cmd.arg("-k").arg("-X").arg("POST"); + if !token.is_empty() { + cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim())); + } + 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") - .arg("-k") - .arg("-X") - .arg("POST") - .arg(format!("https://127.0.0.1:{}/shutdown", port)) - .output(); + let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default(); + let mut cmd = std::process::Command::new("curl"); + cmd.arg("-k").arg("-X").arg("POST"); + if !token.is_empty() { + cmd.arg("-H").arg(format!("Authorization: Bearer {}", token.trim())); + } + 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> { } } - 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