AkurAI-EC2 / crates/platform-mcp/src/server.rs

Latest change 7338926f76e72e7a69a4447469ed20e1488d7989 - refactor(ec2): centralize backup dispatch validation by Ólafur Búi Ólafsson

use crate::tools::ToolRegistry;
use akurai_http::{Method, Request, Response};
use std::process::Command;
use std::sync::Arc;

pub struct McpConfig {
    pub host: String,
    pub port: u16,
    pub token: String,
    pub monitor_url: String,
    pub akurai_ec2_bin: String,
    pub apps_toml: String,
}

pub struct McpServer {
    config: McpConfig,
    tools: Arc<ToolRegistry>,
}

impl McpServer {
    pub fn new(config: McpConfig) -> Result<Self, String> {
        let tools = ToolRegistry::new(
            config.monitor_url.clone(),
            config.akurai_ec2_bin.clone(),
            config.apps_toml.clone(),
        );
        Ok(Self {
            config,
            tools: Arc::new(tools),
        })
    }

    pub fn serve(&self) -> Result<(), String> {
        let addr = format!("{}:{}", self.config.host, self.config.port);
        println!("akurai-platform-mcp listening on http://{addr}/mcp");

        let token = self.config.token.clone();
        let tools = Arc::clone(&self.tools);
        let handler = move |req: &Request| -> Response { handle_http(&token, &tools, req) };

        let server =
            akurai_http::Server::bind(&addr).map_err(|e| format!("failed to bind: {e}"))?;
        server
            .run(handler)
            .map_err(|e| format!("server error: {e}"))
    }
}

fn handle_http(token: &str, tools: &ToolRegistry, req: &Request) -> Response {
    if req.path != "/mcp" {
        return Response::new(404).with_text("Not found");
    }
    if req.method != Method::Post {
        return Response::new(405).with_text("Method not allowed");
    }

    // Bearer token auth with constant-time comparison
    if !token.is_empty() {
        let authorized = req
            .header("Authorization")
            .and_then(|s| s.strip_prefix("Bearer "))
            .map(|t| {
                let supplied = t.trim().as_bytes();
                let expected = token.as_bytes();
                constant_time_equal(supplied, expected)
            })
            .unwrap_or(false);
        if !authorized {
            return Response::new(401).with_text("unauthorized");
        }
    }

    let result = tools.handle_request(req.body_str().as_ref());
    if result.is_empty() {
        return Response::new(202).with_header("Access-Control-Allow-Origin", "*");
    }

    Response::ok()
        .with_header("Access-Control-Allow-Origin", "*")
        .with_header("Content-Type", "application/json")
        .with_body("application/json", result.into_bytes())
}

/// Constant-time byte slice comparison. Returns true if the two slices are
/// equal; the comparison runs in time proportional to `left.len()`, not the
/// position of the first difference — no early exit on mismatch.
fn constant_time_equal(left: &[u8], right: &[u8]) -> bool {
    if left.len() != right.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for (a, b) in left.iter().zip(right.iter()) {
        diff |= a ^ b;
    }
    diff == 0
}

/// Run a command and return (stdout, stderr, exit_code).
pub fn run_command(bin: &str, args: &[&str]) -> Result<(String, String, i32), String> {
    let output = Command::new(bin)
        .args(args)
        .output()
        .map_err(|e| format!("failed to execute {bin}: {e}"))?;
    let stdout = String::from_utf8_lossy(&output.stdout).into_owned();
    let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
    let code = output.status.code().unwrap_or(-1);
    Ok((stdout, stderr, code))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_constant_time_equal_matches() {
        assert!(constant_time_equal(b"abc", b"abc"));
    }

    #[test]
    fn test_constant_time_equal_mismatch_length() {
        assert!(!constant_time_equal(b"abc", b"ab"));
    }

    #[test]
    fn test_constant_time_equal_mismatch_content() {
        assert!(!constant_time_equal(b"abc", b"abd"));
    }

    #[test]
    fn test_constant_time_equal_empty() {
        assert!(constant_time_equal(b"", b""));
    }

    #[test]
    fn test_constant_time_equal_single_byte() {
        assert!(constant_time_equal(b"x", b"x"));
        assert!(!constant_time_equal(b"x", b"y"));
    }

    #[test]
    fn test_constant_time_equal_case_sensitive() {
        // Bearer tokens are case-sensitive
        assert!(!constant_time_equal(b"Token", b"token"));
    }

    // ── HTTP bearer auth tests ───────────────────────────────────

    fn make_request(path: &str, body: &str, auth_header: Option<&str>) -> Request {
        let mut headers = Vec::new();
        if let Some(h) = auth_header {
            headers.push(("Authorization".to_string(), h.to_string()));
        }
        headers.push(("Content-Type".to_string(), "application/json".to_string()));
        Request {
            method: Method::Post,
            path: path.to_string(),
            query: None,
            version: "HTTP/1.1".to_string(),
            headers,
            body: body.as_bytes().to_vec(),
            peer: None,
        }
    }

    fn make_registry() -> ToolRegistry {
        ToolRegistry::new(
            "http://localhost:8800".into(),
            "akurai-ec2".into(),
            "/tmp/nonexistent.toml".into(),
        )
    }

    #[test]
    fn test_bearer_auth_missing_token_rejected() {
        let token = "secret123";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            None, // No Authorization header
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 401);
        assert!(String::from_utf8_lossy(&response.body).contains("unauthorized"));
    }

    #[test]
    fn test_bearer_auth_wrong_token_rejected() {
        let token = "secret123";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            Some("Bearer wrong-token"),
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 401);
    }

    #[test]
    fn test_bearer_auth_correct_token_accepted() {
        let token = "secret123";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            Some("Bearer secret123"),
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 200);
        let body = String::from_utf8_lossy(&response.body);
        assert!(body.contains("ec2_fleet_status"));
    }

    #[test]
    fn test_bearer_auth_empty_token_skips_check() {
        // When token is empty, all requests pass through
        let token = "";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            None,
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 200);
    }

    #[test]
    fn test_bearer_auth_case_sensitive_token() {
        let token = "SecretToken";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            Some("Bearer secrettoken"), // lowercase — should fail
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 401);
    }

    #[test]
    fn test_wrong_path_returns_404() {
        let token = "";
        let reg = make_registry();
        let req = make_request("/wrong", "{}", None);
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 404);
    }

    #[test]
    fn test_get_method_returns_405() {
        let token = "";
        let reg = make_registry();
        let mut req = make_request("/mcp", "", None);
        req.method = Method::Get;
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 405);
    }

    #[test]
    fn test_bearer_auth_with_extra_whitespace() {
        // Token with leading/trailing spaces in the header value
        let token = "secret123";
        let reg = make_registry();
        let req = make_request(
            "/mcp",
            r#"{"jsonrpc":"2.0","id":1,"method":"tools/list"}"#,
            Some("Bearer  secret123  "), // extra spaces — should still match after trim
        );
        let response = handle_http(token, &reg, &req);
        assert_eq!(response.status, 200);
    }
}