oauth_cognito_dynamodb_mcp_server/

oauth_cognito_dynamodb_mcp_server.rs

//! OAuth 2.0 MCP Server with AWS Cognito authentication and DynamoDB storage
//!
//! This example demonstrates a complete MCP (Model Context Protocol) server implementation
//! using AWS Cognito for OAuth 2.0 authentication and DynamoDB for persistent storage.
//! The server follows microkernel architecture principles where independent handlers
//! are composed at runtime.
//!
//! # Features
//! - AWS Cognito OAuth 2.0 authentication
//! - DynamoDB persistent storage for OAuth tokens and client data
//! - MCP over HTTP (streamable)
//! - MCP over Server-Sent Events (SSE)
//! - Microkernel architecture with independent, composable handlers
//!
//! # Required Environment Variables
//! ## Cognito Configuration
//! - `COGNITO_CLIENT_ID`: Your Cognito app client ID
//! - `COGNITO_CLIENT_SECRET`: Your Cognito app client secret (optional for public clients)
//! - `COGNITO_DOMAIN`: Your Cognito domain (e.g., mydomain.auth.us-east-1.amazoncognito.com)
//! - `COGNITO_REGION`: AWS region (e.g., us-east-1)
//! - `COGNITO_USER_POOL_ID`: Your Cognito user pool ID (e.g., us-east-1_XXXXXXXXX)
//! - `COGNITO_SCOPE`: OAuth scopes (default: 'openid email profile phone')
//!
//! ## AWS Configuration (for DynamoDB)
//! - `AWS_ACCESS_KEY_ID`: Your AWS access key ID
//! - `AWS_SECRET_ACCESS_KEY`: Your AWS secret access key
//! - `AWS_REGION`: AWS region (should match COGNITO_REGION)
//!
//! ## Server Configuration
//! - `MCP_HOST`: Server host (default: localhost)
//! - `MCP_PORT`: Server port (default: 8080)
//! - `DYNAMODB_TABLE_NAME`: DynamoDB table name (default: oauth-storage)
//! - `DYNAMODB_CREATE_TABLE`: Whether to auto-create table (default: true)
//!
//! # Usage
//! ```bash
//! # Set environment variables
//! export COGNITO_CLIENT_ID="your_client_id"
//! export COGNITO_CLIENT_SECRET="your_client_secret"
//! export COGNITO_DOMAIN="mydomain.auth.us-east-1.amazoncognito.com"
//! export COGNITO_REGION="us-east-1"
//! export COGNITO_USER_POOL_ID="us-east-1_XXXXXXXXX"
//! export AWS_ACCESS_KEY_ID="your_aws_access_key"
//! export AWS_SECRET_ACCESS_KEY="your_aws_secret_key"
//! export AWS_REGION="us-east-1"
//!
//! # Run the server
//! cargo run --example oauth_cognito_dynamodb_mcp_server
//! ```

use oauth_provider_rs::CognitoOAuthConfig;
use remote_mcp_kernel::{
    config::Config, error::AppResult, microkernel::create_full_cognito_microkernel_dynamodb,
};
use std::env;
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};

#[tokio::main]
async fn main() -> AppResult<()> {
    // Load environment variables
    dotenv::dotenv().ok();

    // Load configuration
    let config = Config::from_env()?;

    // Initialize tracing
    init_tracing(&config)?;

    tracing::info!("Starting MCP OAuth server with Cognito and DynamoDB storage...");

    // Create Cognito OAuth configuration
    let cognito_config = CognitoOAuthConfig {
        client_id: config.cognito.client_id.clone(),
        client_secret: config.cognito.client_secret.clone().unwrap_or_default(),
        redirect_uri: format!(
            "http://{}:{}/oauth/callback",
            config.server.host, config.server.port
        ),
        scope: config.cognito.scope.clone(),
        provider_name: "cognito".to_string(),
    };

    // Get DynamoDB configuration
    let table_name =
        env::var("DYNAMODB_TABLE_NAME").unwrap_or_else(|_| "oauth-storage".to_string());
    let create_table = env::var("DYNAMODB_CREATE_TABLE")
        .unwrap_or_else(|_| "true".to_string())
        .parse::<bool>()
        .unwrap_or(true);

    // Log configuration
    log_startup_info(&config, &table_name, create_table);

    // Create microkernel server with Cognito OAuth and DynamoDB storage
    let microkernel = create_full_cognito_microkernel_dynamodb(
        cognito_config,
        config.cognito.cognito_domain.clone(),
        config.cognito.region.clone(),
        config.cognito.user_pool_id.clone(),
        table_name,
        create_table,
    )
    .await
    .map_err(|e| {
        remote_mcp_kernel::error::AppError::Internal(format!("Failed to create microkernel: {}", e))
    })?;

    // Start the microkernel server
    let bind_address = config.bind_socket_addr()?;
    microkernel.serve(bind_address).await?;

    Ok(())
}

fn init_tracing(config: &Config) -> AppResult<()> {
    tracing_subscriber::registry()
        .with(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| config.logging.level.as_str().into()),
        )
        .with(tracing_subscriber::fmt::layer())
        .init();

    Ok(())
}

fn log_startup_info(config: &Config, table_name: &str, create_table: bool) {
    println!("🚀 Starting MCP OAuth server with Cognito and DynamoDB storage...");
    println!("📋 Configuration:");
    println!("  - Architecture: Microkernel (independent handlers)");
    println!("  - OAuth Provider: AWS Cognito");
    println!("  - Storage Backend: DynamoDB");
    println!("  - Server: {}:{}", config.server.host, config.server.port);
    println!("  - Version: {}", config.server.version);
    println!();

    println!("🔐 AWS Cognito Configuration:");
    println!(
        "  - Client ID: {}",
        if config.cognito.client_id.is_empty() {
            "Not configured"
        } else {
            "Configured"
        }
    );
    println!(
        "  - Client Secret: {}",
        match &config.cognito.client_secret {
            Some(secret) if !secret.is_empty() => "Configured",
            _ => "Not configured (Public Client)",
        }
    );
    println!(
        "  - Domain: {}",
        if config.cognito.cognito_domain.is_empty() {
            "Not configured"
        } else {
            &config.cognito.cognito_domain
        }
    );
    println!(
        "  - Region: {}",
        if config.cognito.region.is_empty() {
            "Not configured"
        } else {
            &config.cognito.region
        }
    );
    println!(
        "  - User Pool ID: {}",
        if config.cognito.user_pool_id.is_empty() {
            "Not configured"
        } else {
            &config.cognito.user_pool_id
        }
    );
    println!("  - Scopes: {}", config.cognito.scope);
    println!();

    println!("🗄️  DynamoDB Storage Configuration:");
    println!("  - Table Name: {}", table_name);
    println!("  - Auto-create Table: {}", create_table);
    println!("  - TTL Attribute: expires_at");
    println!();

    println!("🔧 Handlers:");
    println!("  - OAuth Provider (Cognito authentication & authorization)");
    println!("  - Streamable HTTP Handler (MCP over HTTP)");
    println!("  - SSE Handler (MCP over SSE)");
    println!();

    println!("🔐 Required Environment Variables:");
    println!("  ## Cognito Configuration");
    println!("  - COGNITO_CLIENT_ID: Your Cognito app client ID");
    println!(
        "  - COGNITO_CLIENT_SECRET: Your Cognito app client secret (optional for public clients)"
    );
    println!(
        "  - COGNITO_DOMAIN: Your Cognito domain (e.g., mydomain.auth.us-east-1.amazoncognito.com)"
    );
    println!("  - COGNITO_REGION: AWS region (e.g., us-east-1)");
    println!("  - COGNITO_USER_POOL_ID: Your Cognito user pool ID (e.g., us-east-1_XXXXXXXXX)");
    println!("  - COGNITO_SCOPE: OAuth scopes (default: 'openid email profile phone')");
    println!();
    println!("  ## AWS Configuration (for DynamoDB)");
    println!("  - AWS_ACCESS_KEY_ID: Your AWS access key ID");
    println!("  - AWS_SECRET_ACCESS_KEY: Your AWS secret access key");
    println!("  - AWS_REGION: AWS region (should match COGNITO_REGION)");
    println!();
    println!("  ## Server Configuration");
    println!("  - MCP_HOST: Server host (default: localhost)");
    println!("  - MCP_PORT: Server port (default: 8080)");
    println!("  - DYNAMODB_TABLE_NAME: DynamoDB table name (default: oauth-storage)");
    println!("  - DYNAMODB_CREATE_TABLE: Whether to auto-create table (default: true)");
    println!();

    println!("🌐 OAuth 2.0 Endpoints:");
    println!(
        "  - Authorization: https://{}/oauth2/authorize",
        config.cognito.cognito_domain
    );
    println!(
        "  - Token: https://{}/oauth2/token",
        config.cognito.cognito_domain
    );
    println!(
        "  - JWKS: https://{}/oauth2/jwks",
        config.cognito.cognito_domain
    );
    println!(
        "  - UserInfo: https://{}/oauth2/userInfo",
        config.cognito.cognito_domain
    );
    println!();

    println!("💾 DynamoDB Table Configuration:");
    println!("  - Table Structure: Single table design with 'pk' primary key");
    println!("  - TTL Support: Automatic cleanup using 'expires_at' attribute");
    println!("  - Billing Mode: Pay-per-request (on-demand)");
    println!("  - Encryption: Server-side encryption enabled");
    println!();

    println!("🏗️  Microkernel Architecture:");
    println!("  - Independent handlers that can operate standalone");
    println!("  - Runtime composition of services");
    println!("  - Single responsibility per handler");
    println!("  - Easy testing and maintenance");
    println!();
}