server.rs

//! server entry point
//!
//! This module provides the HTTP/HTTPS server implementation for the PDS,
//! using Axum as the web framework.

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

use axum::{
    Router,
    body::Body,
    extract::{ConnectInfo, DefaultBodyLimit, Request, State},
    middleware::{self, Next},
    response::Response,
};
use tokio::net::TcpListener;
use tower_cookies::CookieManagerLayer;
use tower_http::cors::{Any, CorsLayer};

use super::admin;
use super::background_jobs::BackgroundJobs;
use super::db::{PdsDb};
use crate::pds::xrpc::get_caller_info;
use super::oauth;
use super::spaces;
use super::xrpc;
use crate::fs::LocalFileSystem;
use crate::log::Logger;

/// Shared state for the PDS server.
///
/// This is passed to all handlers via Axum's state extraction.
pub struct PdsState {
    /// Logger instance.
    pub log: &'static Logger,
    /// Local file system access.
    pub lfs: LocalFileSystem,
    /// PDS database access.
    pub db: PdsDb,
}

/// PDS Runner - runs the Personal Data Server HTTP endpoints.
pub struct PdsRunner {
    /// Shared state for all handlers.
    state: Arc<PdsState>,
    /// Listen scheme (http or https).
    listen_scheme: String,
    /// Listen host.
    listen_host: String,
    /// Listen port.
    listen_port: i32,
}

impl PdsRunner {
    /// Initialize a new PDS runner.
    ///
    /// Loads configuration from the database and prepares the server for running.
    ///
    /// # Arguments
    ///
    /// * `lfs` - LocalFileSystem instance
    /// * `log` - Logger instance reference (static lifetime)
    ///
    /// # Returns
    ///
    /// A PdsRunner instance ready to run, or an error if initialization fails.
    pub fn initialize(lfs: LocalFileSystem, log: &'static Logger) -> Result<Self, PdsRunnerError> {
        // Connect to PDS database
        let db = PdsDb::connect(&lfs)?;

        // Load server configuration
        let listen_scheme = db.get_config_property("ServerListenScheme")?;
        let listen_host = db.get_config_property("ServerListenHost")?;
        let listen_port = db.get_config_property_int("ServerListenPort")?;

        log.info(&format!(
            "PDS server initialized: {}://{}:{}",
            listen_scheme, listen_host, listen_port
        ));

        let state = Arc::new(PdsState { log, lfs, db });

        Ok(Self {
            state,
            listen_scheme,
            listen_host,
            listen_port,
        })
    }

    /// Run the PDS runner.
    ///
    /// This starts the HTTP server and blocks until shutdown.
    pub async fn run(&self) -> Result<(), PdsRunnerError> {
        self.state.log.info("");
        self.state.log.info("!! Running PDS !!");
        self.state.log.info("");
        self.state.log.info(&format!(
            "admin: {}://{}:{}/admin/",
            self.listen_scheme, self.listen_host, self.listen_port
        ));
        self.state.log.info("");

        // Start background jobs
        let bg_db = PdsDb::connect(&self.state.lfs)?;
        let mut background_jobs = BackgroundJobs::new(
            self.state.lfs.clone(),
            self.state.log,
            std::sync::Arc::new(bg_db),
        );
        background_jobs.start();

        // Build the router
        let app = self.build_router();

        // Create the listener
        let bind_addr = format!("{}:{}", self.listen_host, self.listen_port);
        let listener = TcpListener::bind(&bind_addr).await.map_err(|e| {
            PdsRunnerError::IoError(format!("Failed to bind to {}: {}", bind_addr, e))
        })?;

        self.state
            .log
            .info(&format!("Listening on {}", bind_addr));

        // Run the server with ConnectInfo to capture client socket addresses
        axum::serve(
            listener,
            app.into_make_service_with_connect_info::<SocketAddr>(),
        )
            .await
            .map_err(|e| PdsRunnerError::IoError(format!("Server error: {}", e)))?;

        Ok(())
    }

    /// Build the Axum router with all endpoints.
    fn build_router(&self) -> Router {
        // CORS layer - allow any origin for development
        let cors = CorsLayer::new()
            .allow_origin(Any)
            .allow_methods(Any)
            .allow_headers(Any);

        // =================================================================
        // XRPC ROUTES - AT Protocol XRPC endpoints.
        // =================================================================
        // =================================================================
        // ADMIN ROUTES - This is the AUTHORITATIVE location for all admin routes.
        // When adding a new admin page, add the route here (not in admin/mod.rs).
        // See admin/mod.rs for the checklist of steps to add a new admin page.
        // =================================================================
        Router::new()
            // XRPC endpoints
            .route("/hello", axum::routing::get(xrpc::hello))
            .route("/xrpc/_health", axum::routing::get(xrpc::health))
            .route("/xrpc/com.atproto.server.describeServer", axum::routing::get(xrpc::describe_server))
            .route("/xrpc/com.atproto.identity.resolveHandle", axum::routing::get(xrpc::resolve_handle))
            // Authentication endpoints
            .route("/xrpc/com.atproto.server.createSession", axum::routing::post(xrpc::create_session))
            .route("/xrpc/com.atproto.server.getSession", axum::routing::get(xrpc::get_session))
            .route("/xrpc/com.atproto.server.refreshSession", axum::routing::post(xrpc::refresh_session))
            .route("/xrpc/com.atproto.server.getServiceAuth", axum::routing::get(xrpc::get_service_auth))
            .route("/xrpc/com.atproto.server.checkAccountStatus", axum::routing::get(xrpc::check_account_status))
            .route("/xrpc/com.atproto.server.activateAccount", axum::routing::post(xrpc::activate_account))
            .route("/xrpc/com.atproto.server.deactivateAccount", axum::routing::post(xrpc::deactivate_account))
            // Repo operation endpoints
            .route("/xrpc/com.atproto.repo.describeRepo", axum::routing::get(xrpc::describe_repo))
            .route("/xrpc/com.atproto.repo.getRecord", axum::routing::get(xrpc::get_record))
            .route("/xrpc/com.atproto.repo.listRecords", axum::routing::get(xrpc::list_records))
            .route("/xrpc/com.atproto.repo.createRecord", axum::routing::post(xrpc::create_record))
            .route("/xrpc/com.atproto.repo.putRecord", axum::routing::post(xrpc::put_record))
            .route("/xrpc/com.atproto.repo.deleteRecord", axum::routing::post(xrpc::delete_record))
            .route("/xrpc/com.atproto.repo.applyWrites", axum::routing::post(xrpc::apply_writes))
            // Blob endpoints
            .merge(
                Router::new()
                    .route("/xrpc/com.atproto.repo.uploadBlob", axum::routing::post(xrpc::upload_blob))
                    .layer(DefaultBodyLimit::max(50 * 1024 * 1024)) // 50MB max for blob uploads
            )
            .route("/xrpc/com.atproto.sync.listBlobs", axum::routing::get(xrpc::list_blobs))
            .route("/xrpc/com.atproto.sync.getBlob", axum::routing::get(xrpc::get_blob))
            // Sync endpoints
            .route("/xrpc/com.atproto.sync.getRepo", axum::routing::get(xrpc::sync_get_repo))
            .route("/xrpc/com.atproto.sync.getRecord", axum::routing::get(xrpc::sync_get_record))
            .route("/xrpc/com.atproto.sync.listRepos", axum::routing::get(xrpc::sync_list_repos))
            .route("/xrpc/com.atproto.sync.getRepoStatus", axum::routing::get(xrpc::sync_get_repo_status))
            .route("/xrpc/com.atproto.sync.subscribeRepos", axum::routing::get(xrpc::subscribe_repos))
            // Permissioned spaces
            .route("/xrpc/com.atproto.space.getDelegationToken", axum::routing::get(spaces::get_delegation_token))
            .route("/xrpc/com.atproto.space.getSpaceCredential", axum::routing::post(spaces::get_space_credential))
            .route("/xrpc/com.atproto.space.createRecord", axum::routing::post(spaces::create_space_record))
            .route("/xrpc/com.atproto.space.putRecord", axum::routing::post(spaces::put_space_record))
            .route("/xrpc/com.atproto.space.registerNotify", axum::routing::post(spaces::register_notify))
            .route("/xrpc/com.atproto.space.unregisterNotify", axum::routing::post(spaces::unregister_notify))
            // Simplespace management
            .route("/xrpc/com.atproto.simplespace.createSpace", axum::routing::post(spaces::create_space))
            .route("/xrpc/com.atproto.simplespace.getSpace", axum::routing::get(spaces::get_space))
            // App.bsky endpoints (preferences are handled locally, others proxy to AppView)
            .route("/xrpc/app.bsky.actor.getPreferences", axum::routing::get(xrpc::get_preferences))
            .route("/xrpc/app.bsky.actor.putPreferences", axum::routing::post(xrpc::put_preferences))
            // Catch-all for app.bsky.* and chat.bsky.* routes - proxy to AppView
            .fallback(xrpc::app_bsky_fallback)
            // Static file endpoints
            .route("/", axum::routing::get(xrpc::root))
            .route("/favicon.ico", axum::routing::get(xrpc::favicon))
            // .well-known endpoints
            .route("/.well-known/did.json", axum::routing::get(xrpc::well_known_did))
            .route("/.well-known/atproto-did", axum::routing::get(xrpc::well_known_atproto_did))
            // OAuth endpoints
            .route("/.well-known/oauth-protected-resource", axum::routing::get(oauth::oauth_protected_resource))
            .route("/.well-known/oauth-authorization-server", axum::routing::get(oauth::oauth_authorization_server))
            .route("/oauth/jwks", axum::routing::get(oauth::oauth_jwks))
            .route("/oauth/par", axum::routing::post(oauth::oauth_par))
            .route("/oauth/authorize", axum::routing::get(oauth::oauth_authorize_get).post(oauth::oauth_authorize_post))
            .route("/oauth/token", axum::routing::post(oauth::oauth_token))
            .route("/oauth/revoke", axum::routing::post(oauth::oauth_revoke))
            .route("/oauth/passkeyauthenticationoptions", axum::routing::post(oauth::passkey_authentication_options))
            .route("/oauth/authenticatepasskey", axum::routing::post(oauth::authenticate_passkey))
            // Admin endpoints
            .route("/admin", axum::routing::get(admin::admin_home))
            .route("/admin/", axum::routing::get(admin::admin_home))
            .route("/admin/login", axum::routing::get(admin::admin_login_get).post(admin::admin_login_post))
            .route("/admin/login/", axum::routing::get(admin::admin_login_get).post(admin::admin_login_post))
            .route("/admin/passkeyauthenticationoptions", axum::routing::post(admin::admin_passkey_authentication_options))
            .route("/admin/authenticatepasskey", axum::routing::post(admin::admin_authenticate_passkey))
            .route("/admin/register-passkey", axum::routing::get(admin::admin_register_passkey_get))
            .route("/admin/passkeyregistrationoptions", axum::routing::post(admin::admin_passkey_registration_options))
            .route("/admin/registerpasskey", axum::routing::post(admin::admin_register_passkey_post))
            .route("/admin/logout", axum::routing::post(admin::admin_logout))
            .route("/admin/sessions", axum::routing::get(admin::admin_sessions))
            .route("/admin/sessions/", axum::routing::get(admin::admin_sessions))
            .route("/admin/deletelegacysession", axum::routing::post(admin::admin_delete_legacy_session))
            .route("/admin/deleteoauthsession", axum::routing::post(admin::admin_delete_oauth_session))
            .route("/admin/deleteadminsession", axum::routing::post(admin::admin_delete_admin_session))
            .route("/admin/stats", axum::routing::get(admin::admin_stats))
            .route("/admin/stats/", axum::routing::get(admin::admin_stats))
            .route("/admin/stats_writes", axum::routing::get(admin::admin_stats_writes))
            .route("/admin/stats_writes/", axum::routing::get(admin::admin_stats_writes))
            .route("/admin/stats_ip", axum::routing::get(admin::admin_stats_ip))
            .route("/admin/stats_ip/", axum::routing::get(admin::admin_stats_ip))
            .route("/admin/spaces", axum::routing::get(admin::admin_spaces))
            .route("/admin/spaces/", axum::routing::get(admin::admin_spaces))
            .route("/admin/deletespace", axum::routing::post(admin::admin_delete_space))
            .route("/admin/deletenotifyregistration", axum::routing::post(admin::admin_delete_notify_registration))
            .route("/admin/deletestatistic", axum::routing::post(admin::admin_delete_statistic))
            .route("/admin/deleteallstatistics", axum::routing::post(admin::admin_delete_all_statistics))
            .route("/admin/deleteoldstatistics", axum::routing::post(admin::admin_delete_old_statistics))
            .route("/admin/passkeys", axum::routing::get(admin::admin_passkeys))
            .route("/admin/passkeys/", axum::routing::get(admin::admin_passkeys))
            .route("/admin/deletepasskey", axum::routing::post(admin::admin_delete_passkey))
            .route("/admin/deletepasskeychallenge", axum::routing::post(admin::admin_delete_passkey_challenge))
            .route("/admin/config", axum::routing::get(admin::admin_config_get).post(admin::admin_config_post))
            .route("/admin/config/", axum::routing::get(admin::admin_config_get).post(admin::admin_config_post))
            .route("/admin/actions", axum::routing::get(admin::admin_actions_get).post(admin::admin_actions_post))
            .route("/admin/actions/", axum::routing::get(admin::admin_actions_get).post(admin::admin_actions_post))
            .layer(middleware::from_fn_with_state(
                self.state.clone(),
                logging_middleware,
            ))
            .layer(CookieManagerLayer::new())
            .layer(cors)
            .with_state(self.state.clone())
    }
}

/// Logging middleware that logs all HTTP requests and responses.
async fn logging_middleware(
    State(state): State<Arc<PdsState>>,
    ConnectInfo(addr): ConnectInfo<SocketAddr>,
    request: Request,
    next: Next,
) -> Response {
    let method = request.method().clone();
    let uri = request.uri().clone();
    let path = uri.path().to_string();
    let start = std::time::Instant::now();

    // Extract caller info for statistics
    let (ip_address, user_agent) = get_caller_info(request.headers(), Some(addr));

    // get params string for logging
    let mut params_str = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();

    // this one is verbose and we'll never read it
    if path == "/xrpc/app.bsky.feed.getPosts"
    {
        params_str = "<params_str omitted>".to_string();
    }

    // Log the connection
    state.log.info(&format!(
        "[BEGIN REQUEST] {} {} {} {} {}",
        ip_address, method, path, params_str, user_agent,
    ));

    // Optionally log the full details of configured XRPC endpoints (debugging
    // aid controlled by the LogXrpcEndpoints config property).
    let request = maybe_log_xrpc_request(&state, &uri, &path, request).await;

    // Run the next handler
    let response = next.run(request).await;

    let elapsed = start.elapsed();
    let status = response.status();

    state.log.info(&format!(
        "[END REQUEST] [{} ({:.2?})] {} {} {} {} {}",
        status.as_u16(), elapsed, ip_address, method, path, params_str, user_agent
    ));

    response
}

/// Maximum request body size (bytes) buffered for XRPC request logging.
const MAX_LOGGED_XRPC_REQUEST_BYTES: usize = 50 * 1024 * 1024;

/// If the request targets an XRPC endpoint configured in `LogXrpcEndpoints`,
/// buffer its body, log the full request at info level, and return a
/// reconstructed request so downstream handlers still see the body.
///
/// For all other requests the original request is returned untouched.
async fn maybe_log_xrpc_request(
    state: &Arc<PdsState>,
    uri: &axum::http::Uri,
    path: &str,
    request: Request,
) -> Request {
    let nsid = match path.strip_prefix("/xrpc/") {
        Some(nsid) => nsid,
        None => return request,
    };

    let endpoints = match state.db.get_config_property_hash_set("LogXrpcEndpoints") {
        Ok(endpoints) => endpoints,
        Err(_) => return request,
    };

    if !endpoints.contains(nsid) {
        return request;
    }

    let method = request.method().clone();
    let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
    let headers = request.headers().clone();

    let (parts, body) = request.into_parts();
    match axum::body::to_bytes(body, MAX_LOGGED_XRPC_REQUEST_BYTES).await {
        Ok(bytes) => {
            xrpc::log_xrpc_request(state, &method, path, &query, &headers, &bytes);
            Request::from_parts(parts, Body::from(bytes))
        }
        Err(_) => {
            state.log.warning(&format!(
                "[LOG_XRPC] failed to buffer request body for {}", path
            ));
            Request::from_parts(parts, Body::empty())
        }
    }
}
#[derive(thiserror::Error, Debug)]
pub enum PdsRunnerError {
    #[error("Database error: {0}")]
    DbError(#[from] super::db::PdsDbError),

    #[error("IO error: {0}")]
    IoError(String),
}