secrets/crates/secrets-mcp/src/main.rs

mod auth;
mod error;
mod logging;
mod oauth;
mod tools;
mod web;

use std::net::SocketAddr;
use std::sync::Arc;

use anyhow::{Context, Result};
use axum::Router;
use rmcp::transport::streamable_http_server::{
    StreamableHttpService, session::local::LocalSessionManager,
};
use sqlx::PgPool;
use tower_http::cors::{Any, CorsLayer};
use tower_sessions::cookie::SameSite;
use tower_sessions::session_store::ExpiredDeletion;
use tower_sessions::{Expiry, SessionManagerLayer};
use tower_sessions_sqlx_store_chrono::PostgresStore;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::fmt::time::FormatTime;

use secrets_core::config::resolve_db_config;
use secrets_core::db::{create_pool, migrate};

use crate::oauth::OAuthConfig;
use crate::tools::SecretsService;

/// Shared application state injected into web routes and middleware.
#[derive(Clone)]
pub struct AppState {
    pub pool: PgPool,
    pub google_config: Option<OAuthConfig>,
    pub base_url: String,
    pub http_client: reqwest::Client,
}

fn load_env_var(name: &str) -> Option<String> {
    std::env::var(name).ok().filter(|s| !s.is_empty())
}

/// Pretty-print bind address in logs (`127.0.0.1` → `localhost`); actual socket bind unchanged.
fn listen_addr_log_display(bind_addr: &str) -> String {
    bind_addr
        .strip_prefix("127.0.0.1:")
        .map(|port| format!("localhost:{port}"))
        .unwrap_or_else(|| bind_addr.to_string())
}

fn load_oauth_config(prefix: &str, base_url: &str, path: &str) -> Option<OAuthConfig> {
    let client_id = load_env_var(&format!("{}_CLIENT_ID", prefix))?;
    let client_secret = load_env_var(&format!("{}_CLIENT_SECRET", prefix))?;
    Some(OAuthConfig {
        client_id,
        client_secret,
        redirect_uri: format!("{}{}", base_url, path),
    })
}

/// Log line timestamps in the process local timezone (honors `TZ` / system zone).
#[derive(Clone, Copy, Default)]
struct LocalRfc3339Time;

impl FormatTime for LocalRfc3339Time {
    fn format_time(&self, w: &mut tracing_subscriber::fmt::format::Writer<'_>) -> std::fmt::Result {
        write!(
            w,
            "{}",
            chrono::Local::now().to_rfc3339_opts(chrono::SecondsFormat::Millis, false)
        )
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // Load .env if present
    let _ = dotenvy::dotenv();

    tracing_subscriber::fmt()
        .with_timer(LocalRfc3339Time)
        .with_env_filter(
            EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "secrets_mcp=info,tower_http=info".into()),
        )
        .init();

    // ── Database ──────────────────────────────────────────────────────────────
    let db_config = resolve_db_config("")
        .context("Database not configured. Set SECRETS_DATABASE_URL environment variable.")?;
    let pool = create_pool(&db_config)
        .await
        .context("failed to connect to database")?;
    migrate(&pool)
        .await
        .context("failed to run database migrations")?;
    tracing::info!("Database connected and migrated");

    // ── Configuration ─────────────────────────────────────────────────────────
    let base_url = load_env_var("BASE_URL").unwrap_or_else(|| "http://localhost:9315".to_string());
    let bind_addr =
        load_env_var("SECRETS_MCP_BIND").unwrap_or_else(|| "127.0.0.1:9315".to_string());

    // ── OAuth providers ───────────────────────────────────────────────────────
    let google_config = load_oauth_config("GOOGLE", &base_url, "/auth/google/callback");

    if google_config.is_none() {
        tracing::warn!(
            "No OAuth providers configured. Set GOOGLE_CLIENT_ID/GOOGLE_CLIENT_SECRET to enable login."
        );
    }

    // ── Session store (PostgreSQL-backed) ─────────────────────────────────────
    let session_store = PostgresStore::new(pool.clone());
    session_store
        .migrate()
        .await
        .context("failed to run session table migration")?;
    // Prune expired rows every hour; task is aborted when the server shuts down.
    let session_cleanup = tokio::spawn(
        session_store
            .clone()
            .continuously_delete_expired(tokio::time::Duration::from_secs(3600)),
    );
    // Strict would drop the session cookie on redirect from Google → our origin (cross-site nav).
    let session_layer = SessionManagerLayer::new(session_store)
        .with_secure(base_url.starts_with("https://"))
        .with_same_site(SameSite::Lax)
        .with_expiry(Expiry::OnInactivity(time::Duration::days(14)));

    // ── App state ─────────────────────────────────────────────────────────────
    let app_state = AppState {
        pool: pool.clone(),
        google_config,
        base_url: base_url.clone(),
        http_client: reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(15))
            .build()
            .context("failed to build HTTP client")?,
    };

    // ── MCP service ───────────────────────────────────────────────────────────
    let pool_arc = Arc::new(pool.clone());

    let mcp_service = StreamableHttpService::new(
        move || {
            let p = pool_arc.clone();
            Ok(SecretsService::new(p))
        },
        LocalSessionManager::default().into(),
        Default::default(),
    );

    // ── Router ────────────────────────────────────────────────────────────────
    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods(Any)
        .allow_headers(Any);

    let router = Router::new()
        .merge(web::web_router())
        .nest_service("/mcp", mcp_service)
        .layer(axum::middleware::from_fn(
            logging::request_logging_middleware,
        ))
        .layer(axum::middleware::from_fn_with_state(
            pool,
            auth::bearer_auth_middleware,
        ))
        .layer(session_layer)
        .layer(cors)
        .with_state(app_state);

    // ── Start server ──────────────────────────────────────────────────────────
    let listener = tokio::net::TcpListener::bind(&bind_addr)
        .await
        .with_context(|| format!("failed to bind to {}", bind_addr))?;

    tracing::info!(
        "Secrets MCP Server listening on http://{}",
        listen_addr_log_display(&bind_addr)
    );
    tracing::info!("MCP endpoint: {}/mcp", base_url);

    axum::serve(
        listener,
        router.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown_signal())
    .await
    .context("server error")?;

    session_cleanup.abort();
    Ok(())
}

async fn shutdown_signal() {
    tokio::signal::ctrl_c()
        .await
        .expect("failed to install CTRL+C signal handler");
    tracing::info!("Shutting down gracefully...");
}