photoncloud-monorepo/iam/crates/iam-server/src/main.rs

//! IAM Server
//!
//! The main entry point for the IAM gRPC server.

mod config;
mod rest;

use std::sync::Arc;
use std::time::Duration;

use clap::Parser;
use metrics_exporter_prometheus::PrometheusBuilder;
use tonic::transport::{Certificate, Identity, Server, ServerTlsConfig};
use tonic_health::server::health_reporter;
use tracing::{info, warn};

use iam_api::{
    iam_admin_server::IamAdminServer, iam_authz_server::IamAuthzServer,
    iam_token_server::IamTokenServer, IamAdminService, IamAuthzService, IamTokenService,
};
use iam_authn::{InternalTokenConfig, InternalTokenService, SigningKey};
use iam_authz::{PolicyCache, PolicyCacheConfig, PolicyEvaluator};
use iam_store::{Backend, BackendConfig, BindingStore, PrincipalStore, RoleStore, TokenStore};

use config::{BackendKind, ServerConfig};

/// IAM Server
#[derive(Parser, Debug)]
#[command(name = "iam-server")]
#[command(about = "Identity and Access Management Server")]
struct Args {
    /// Configuration file path
    #[arg(short, long)]
    config: Option<String>,

    /// Listen address (overrides config)
    #[arg(long)]
    addr: Option<String>,

    /// Log level (overrides config)
    #[arg(long)]
    log_level: Option<String>,

    /// Metrics port for Prometheus scraping (default: 9093)
    #[arg(long, default_value = "9093")]
    metrics_port: u16,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = Args::parse();

    // Load configuration
    let mut config = match &args.config {
        Some(path) => ServerConfig::from_file(path)?,
        None => ServerConfig::from_env()?,
    };

    // Apply CLI overrides
    if let Some(addr) = args.addr {
        config.server.addr = addr.parse()?;
    }
    if let Some(level) = args.log_level {
        config.logging.level = level;
    }

    // Initialize logging
    init_logging(&config.logging.level);

    // Initialize Prometheus metrics exporter
    // Serves metrics at http://0.0.0.0:{metrics_port}/metrics
    let metrics_addr = format!("0.0.0.0:{}", args.metrics_port);
    let builder = PrometheusBuilder::new();
    builder
        .with_http_listener(metrics_addr.parse::<std::net::SocketAddr>()?)
        .install()
        .expect("Failed to install Prometheus metrics exporter");

    info!(
        "Prometheus metrics available at http://{}/metrics",
        metrics_addr
    );

    // Register common metrics
    metrics::describe_counter!(
        "iam_authz_requests_total",
        "Total number of authorization requests"
    );
    metrics::describe_counter!(
        "iam_authz_allowed_total",
        "Total number of allowed authorization requests"
    );
    metrics::describe_counter!(
        "iam_authz_denied_total",
        "Total number of denied authorization requests"
    );
    metrics::describe_counter!("iam_token_issued_total", "Total number of tokens issued");
    metrics::describe_histogram!(
        "iam_request_duration_seconds",
        "Request duration in seconds"
    );

    info!("Starting IAM server on {}", config.server.addr);

    // Create backend
    let backend = create_backend(&config.store).await?;
    let backend = Arc::new(backend);

    // Create stores
    let principal_store = Arc::new(PrincipalStore::new(backend.clone()));
    let role_store = Arc::new(RoleStore::new(backend.clone()));
    let binding_store = Arc::new(BindingStore::new(backend.clone()));
    let token_store = Arc::new(TokenStore::new(backend.clone()));

    // Initialize builtin roles
    info!("Initializing builtin roles...");
    role_store.init_builtin_roles().await?;

    // Create policy cache
    let cache_config = PolicyCacheConfig {
        binding_ttl: Duration::from_secs(300),
        role_ttl: Duration::from_secs(600),
        max_binding_entries: 10000,
        max_role_entries: 1000,
    };
    let cache = Arc::new(PolicyCache::new(cache_config));

    // Create evaluator
    let evaluator = Arc::new(PolicyEvaluator::new(
        binding_store.clone(),
        role_store.clone(),
        cache,
    ));

    // Create token service
    let signing_key = if config.authn.internal_token.signing_key.is_empty() {
        warn!("No signing key configured, generating random key");
        SigningKey::generate("iam-key-1")
    } else {
        SigningKey::new(
            "iam-key-1",
            config.authn.internal_token.signing_key.as_bytes().to_vec(),
        )
    };

    let token_config = InternalTokenConfig::new(signing_key, &config.authn.internal_token.issuer)
        .with_default_ttl(Duration::from_secs(
            config.authn.internal_token.default_ttl_seconds,
        ))
        .with_max_ttl(Duration::from_secs(
            config.authn.internal_token.max_ttl_seconds,
        ));

    let token_service = Arc::new(InternalTokenService::new(token_config));

    // Create gRPC services
    let authz_service = IamAuthzService::new(evaluator, principal_store.clone());
    let token_grpc_service =
        IamTokenService::new(token_service, principal_store.clone(), token_store.clone());
    let admin_service = IamAdminService::new(
        principal_store.clone(),
        role_store.clone(),
        binding_store.clone(),
    );

    info!("IAM server ready, starting gRPC listeners...");

    // Create health service (for K8s liveness/readiness probes)
    // Uses grpc.health.v1.Health standard protocol
    let (mut health_reporter, health_service) = health_reporter();

    // Mark services as serving
    health_reporter
        .set_serving::<IamAuthzServer<IamAuthzService>>()
        .await;
    health_reporter
        .set_serving::<IamTokenServer<IamTokenService>>()
        .await;
    health_reporter
        .set_serving::<IamAdminServer<IamAdminService>>()
        .await;

    // Spawn health monitoring task
    let backend_for_health = backend.clone();
    tokio::spawn(async move {
        // Periodically check backend connectivity
        loop {
            tokio::time::sleep(Duration::from_secs(30)).await;
            // Backend health check could be added here if Backend exposes a ping method
            let _ = backend_for_health; // Keep reference alive
        }
    });

    info!("Health check service enabled (grpc.health.v1.Health)");

    // Configure TLS if enabled
    let mut server = Server::builder();

    if let Some(tls_config) = &config.server.tls {
        info!("TLS enabled, loading certificates...");
        let cert = tokio::fs::read(&tls_config.cert_file)
            .await
            .map_err(|e| format!("Failed to read cert file: {}", e))?;
        let key = tokio::fs::read(&tls_config.key_file)
            .await
            .map_err(|e| format!("Failed to read key file: {}", e))?;
        let server_identity = Identity::from_pem(cert, key);

        let tls = if tls_config.require_client_cert {
            info!("mTLS enabled, requiring client certificates");
            let ca_cert = tokio::fs::read(
                tls_config
                    .ca_file
                    .as_ref()
                    .ok_or("ca_file required when require_client_cert=true")?,
            )
            .await
            .map_err(|e| format!("Failed to read CA file: {}", e))?;
            let ca = Certificate::from_pem(ca_cert);

            ServerTlsConfig::new()
                .identity(server_identity)
                .client_ca_root(ca)
        } else {
            info!("TLS-only mode, client certificates not required");
            ServerTlsConfig::new().identity(server_identity)
        };

        server = server
            .tls_config(tls)
            .map_err(|e| format!("Failed to configure TLS: {}", e))?;
        info!("TLS configuration applied successfully");
    } else {
        info!("TLS disabled, running in plain-text mode");
    }

    // gRPC server
    let grpc_server = server
        .add_service(health_service)
        .add_service(IamAuthzServer::new(authz_service))
        .add_service(IamTokenServer::new(token_grpc_service))
        .add_service(IamAdminServer::new(admin_service))
        .serve(config.server.addr);

    // HTTP REST API server
    let http_addr = config.server.http_addr;
    let rest_state = rest::RestApiState {
        server_addr: config.server.addr.to_string(),
    };
    let rest_app = rest::build_router(rest_state);
    let http_listener = tokio::net::TcpListener::bind(&http_addr).await?;

    info!(http_addr = %http_addr, "HTTP REST API server starting");

    let http_server = async move {
        axum::serve(http_listener, rest_app)
            .await
            .map_err(|e| format!("HTTP server error: {}", e))
    };

    // Run both servers concurrently
    tokio::select! {
        result = grpc_server => {
            result?;
        }
        result = http_server => {
            result?;
        }
    }

    Ok(())
}

async fn create_backend(
    config: &config::StoreConfig,
) -> Result<Backend, Box<dyn std::error::Error>> {
    let backend_config = match config.backend {
        BackendKind::Memory => {
            info!("Using in-memory backend");
            BackendConfig::Memory
        }
        BackendKind::Chainfire => {
            let endpoints = config
                .chainfire_endpoints
                .clone()
                .ok_or("chainfire_endpoints required for chainfire backend")?;
            info!("Using Chainfire backend with endpoints: {:?}", endpoints);
            BackendConfig::Chainfire { endpoints }
        }
        BackendKind::FlareDb => {
            let endpoint = config
                .flaredb_endpoint
                .clone()
                .ok_or("flaredb_endpoint required for flaredb backend")?;
            let namespace = config
                .flaredb_namespace
                .clone()
                .unwrap_or_else(|| "iam".into());
            info!(
                "Using FlareDB backend at {} (namespace: {})",
                endpoint, namespace
            );
            BackendConfig::FlareDb {
                endpoint,
                namespace,
            }
        }
    };

    Backend::new(backend_config).await.map_err(|e| e.into())
}

fn init_logging(level: &str) {
    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};

    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(level));

    tracing_subscriber::registry()
        .with(filter)
        .with(tracing_subscriber::fmt::layer())
        .init();
}