AkurAI-Notes / crates/notes-mcp/src/server.rs

Latest change 228bc96dddc8e5490cfc2ed27705c0d1dbbfaf83 - MCP: OAuth protected resource for the AkurAI IdP by Ólafur Búi Ólafsson

use crate::oauth::{self, OAuth};
use crate::tools::ToolRegistry;
use akurai_http::{Method, Request, Response};
use akurai_storage::BTree;
use std::sync::{Arc, Mutex};

pub struct McpConfig {
    pub host: String,
    pub port: u16,
    pub db: String,
    pub token: String,
    pub embed_url: Option<String>,
    pub embed_model: Option<String>,
    pub embed_api_key: Option<String>,
    pub owner_sub: String,
}

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

impl McpServer {
    pub fn new(config: McpConfig) -> Result<Self, String> {
        let btree =
            BTree::open(&config.db).map_err(|e| format!("failed to open MCP database: {e}"))?;
        Ok(Self::with_db(config, Arc::new(Mutex::new(btree))))
    }

    pub fn with_db(config: McpConfig, db: Arc<Mutex<BTree>>) -> Self {
        let tools = ToolRegistry::new_background(
            db,
            config.embed_url.clone(),
            config.embed_model.clone(),
            config.embed_api_key.clone(),
            config.owner_sub.clone(),
        );
        Self {
            config,
            tools: Arc::new(tools),
            oauth: None,
        }
    }

    /// Also accept IdP access tokens and publish protected-resource metadata.
    pub fn with_oauth(mut self, oauth: OAuth) -> Self {
        self.oauth = Some(Arc::new(oauth));
        self
    }

    pub fn handle_request(&self, body: &str) -> String {
        self.tools.handle_request(body)
    }

    pub fn handle_http(&self, req: &Request) -> Response {
        handle_http(&self.config.token, self.oauth.as_deref(), &self.tools, req)
    }

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

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

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

/// Paths this server answers, so a host server can route them here.
pub fn owns_path(path: &str) -> bool {
    path == "/mcp" || oauth::is_metadata_path(path)
}

fn handle_http(
    token: &str,
    oauth: Option<&OAuth>,
    tools: &ToolRegistry,
    req: &Request,
) -> Response {
    if let Some(oauth) = oauth.filter(|_| oauth::is_metadata_path(&req.path)) {
        return metadata_response(oauth);
    }
    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");
    }
    if let Some(oauth) = oauth {
        if !oauth.config.origin_allowed(req.header("Origin")) {
            return Response::new(403).with_text("invalid Origin");
        }
    }
    if let Err(challenge) = authorize(token, oauth, req) {
        return challenge;
    }

    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())
}

fn bearer(req: &Request) -> Option<&str> {
    req.header("Authorization")
        .and_then(|s| s.strip_prefix("Bearer "))
        .map(str::trim)
        .filter(|t| !t.is_empty())
}

/// Accept the static MCP token, or (when configured) an IdP access token.
fn authorize(token: &str, oauth: Option<&OAuth>, req: &Request) -> Result<(), Response> {
    let presented = bearer(req);
    if token.is_empty() && oauth.is_none() {
        return Ok(());
    }
    if !token.is_empty() && presented == Some(token) {
        return Ok(());
    }
    let Some(oauth) = oauth else {
        return Err(Response::new(401).with_text("unauthorized"));
    };
    let Some(presented) = presented else {
        return Err(unauthorized(oauth, false));
    };
    oauth
        .authenticate(presented)
        .map(|_| ())
        .map_err(|_| unauthorized(oauth, true))
}

fn unauthorized(oauth: &OAuth, invalid: bool) -> Response {
    Response::new(401)
        .with_header(
            "WWW-Authenticate",
            &oauth::www_authenticate(&oauth.config, invalid),
        )
        .with_text("unauthorized")
}

fn metadata_response(oauth: &OAuth) -> Response {
    Response::ok()
        .with_header("Access-Control-Allow-Origin", "*")
        .with_header("Cache-Control", "public, max-age=3600")
        .with_body(
            "application/json",
            oauth::protected_resource_metadata(&oauth.config)
                .to_string()
                .into_bytes(),
        )
}