pds-gatekeeper / src / main.rs

#![warn(clippy::unwrap_used)]
use crate::oauth_provider::sign_in;
use crate::xrpc::com_atproto_server::{create_account, create_session, get_session, update_email};
use axum::body::Body;
use axum::handler::Handler;
use axum::http::{Method, header};
use axum::middleware as ax_middleware;
use axum::routing::post;
use axum::{Router, routing::get};
use axum_template::engine::Engine;
use handlebars::Handlebars;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::rt::TokioExecutor;
use lettre::{AsyncSmtpTransport, Tokio1Executor};
use rust_embed::RustEmbed;
use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode};
use sqlx::{SqlitePool, sqlite::SqlitePoolOptions};
use std::path::Path;
use std::time::Duration;
use std::{env, net::SocketAddr};
use tower_governor::GovernorLayer;
use tower_governor::governor::GovernorConfigBuilder;
use tower_governor::key_extractor::SmartIpKeyExtractor;
use tower_http::compression::CompressionLayer;
use tower_http::cors::{Any, CorsLayer};
use tracing::log;
use tracing_subscriber::{EnvFilter, fmt, prelude::*};

pub mod helpers;
mod middleware;
mod oauth_provider;
mod xrpc;

type HyperUtilClient = hyper_util::client::legacy::Client<HttpConnector, Body>;

#[derive(RustEmbed)]
#[folder = "email_templates"]
#[include = "*.hbs"]
struct EmailTemplates;

#[derive(Clone)]
pub struct AppState {
    account_pool: SqlitePool,
    pds_gatekeeper_pool: SqlitePool,
    reverse_proxy_client: HyperUtilClient,
    pds_base_url: String,
    mailer: AsyncSmtpTransport<Tokio1Executor>,
    mailer_from: String,
    template_engine: Engine<Handlebars<'static>>,
}

async fn root_handler() -> impl axum::response::IntoResponse {
    let body = r"

        ...oO                _.--X~~OO~~X--._                       ...oOO
                         _.-~   / \ II / \   ~-._
                    [].-~  \   /   \||/   \   /  ~-.[]    ...o
 ...o           _   ||/     \ /     ||     \ /     \||   _
               (_)  |X       X      ||      X       X|  (_)
              _-~-_ ||\     / \     ||     / \     /|| _-~-_
              ||||| || \   /   \   /||\   /   \   / || |||||
              |   |_||  \ /     \ / || \ /     \ /  ||_|   |
              |   |~||   X       X  ||  X       X   ||~|   |
==============|   | ||  / \     / \ || / \     / \  || |   |==============
______________|   | || /   \   /   \||/   \   /   \ || |   |______________
    .     .   |   | ||/     \ /     ||     \ /     \|| |   |  .       .
       /      |   | |X       X      ||      X       X| |   |    /        /
  /   .       |   | ||\     / \     ||     / \     /|| |   | .      /   .
.          /  |   | || \   /   \   /||\   /   \   / || |   |   .  .
    .    .    |   | ||  \ /     \ / || \ /     \ /  || |   |          .
      /       |   | ||   X       X  ||  X       X   || |   | . / .      /
  /        .  |   | ||  / \     / \ || / \     / \  || |   |        /
         /    |   | || /   \   /   \||/   \   /   \ || |   |   .         /
.    .    .   |   | ||/     \ /    /||\    \ /     \|| |   |     /.    .
              |   |_|X       X    / II \    X       X|_|   |  .     .   /
==============|   |~II~~~~~~~~~~~~~~OO~~~~~~~~~~~~~~II~|   |==============
 ";

    let intro = "\n\nThis is a PDS gatekeeper\n\nCode: https://tangled.sh/@baileytownsend.dev/pds-gatekeeper\n";

    let banner = format!("         {body}\n{intro}");

    (
        [(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
        banner,
    )
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    setup_tracing();
    let pds_env_location =
        env::var("PDS_ENV_LOCATION").unwrap_or_else(|_| "/pds/pds.env".to_string());

    let result_of_finding_pds_env = dotenvy::from_path(Path::new(&pds_env_location));
    if let Err(e) = result_of_finding_pds_env {
        log::error!(
            "Error loading pds.env file (ignore if you loaded your variables in the environment somehow else): {e}"
        );
    }

    let pds_root =
        env::var("PDS_DATA_DIRECTORY").expect("PDS_DATA_DIRECTORY is not set in your pds.env file");
    let account_db_url = format!("{pds_root}/account.sqlite");

    let account_options = SqliteConnectOptions::new()
        .filename(account_db_url)
        .busy_timeout(Duration::from_secs(5));

    let account_pool = SqlitePoolOptions::new()
        .max_connections(5)
        .connect_with(account_options)
        .await?;

    let bells_db_url = format!("{pds_root}/pds_gatekeeper.sqlite");
    let options = SqliteConnectOptions::new()
        .journal_mode(SqliteJournalMode::Wal)
        .filename(bells_db_url)
        .create_if_missing(true)
        .busy_timeout(Duration::from_secs(5));
    let pds_gatekeeper_pool = SqlitePoolOptions::new()
        .max_connections(5)
        .connect_with(options)
        .await?;

    // Run migrations for the extra database
    // Note: the migrations are embedded at compile time from the given directory
    // sqlx
    sqlx::migrate!("./migrations")
        .run(&pds_gatekeeper_pool)
        .await?;

    let client: HyperUtilClient =
        hyper_util::client::legacy::Client::<(), ()>::builder(TokioExecutor::new())
            .build(HttpConnector::new());

    //Emailer set up
    let smtp_url =
        env::var("PDS_EMAIL_SMTP_URL").expect("PDS_EMAIL_SMTP_URL is not set in your pds.env file");
    let sent_from = env::var("PDS_EMAIL_FROM_ADDRESS")
        .expect("PDS_EMAIL_FROM_ADDRESS is not set in your pds.env file");

    let mailer: AsyncSmtpTransport<Tokio1Executor> =
        AsyncSmtpTransport::<Tokio1Executor>::from_url(smtp_url.as_str())?.build();
    //Email templates setup
    let mut hbs = Handlebars::new();

    let users_email_directory = env::var("GATEKEEPER_EMAIL_TEMPLATES_DIRECTORY");
    if let Ok(users_email_directory) = users_email_directory {
        hbs.register_template_file(
            "two_factor_code.hbs",
            format!("{users_email_directory}/two_factor_code.hbs"),
        )?;
    } else {
        let _ = hbs.register_embed_templates::<EmailTemplates>();
    }

    let pds_base_url =
        env::var("PDS_BASE_URL").unwrap_or_else(|_| "http://localhost:3000".to_string());

    let state = AppState {
        account_pool,
        pds_gatekeeper_pool,
        reverse_proxy_client: client,
        pds_base_url,
        mailer,
        mailer_from: sent_from,
        template_engine: Engine::from(hbs),
    };

    // Rate limiting
    //Allows 5 within 60 seconds, and after 60 should drop one off? So hit 5, then goes to 4 after 60 seconds.
    let create_session_governor_conf = GovernorConfigBuilder::default()
        .per_second(60)
        .burst_size(5)
        .key_extractor(SmartIpKeyExtractor)
        .finish()
        .expect("failed to create governor config for create session. this should not happen and is a bug");

    // Create a second config with the same settings for the other endpoint
    let sign_in_governor_conf = GovernorConfigBuilder::default()
        .per_second(60)
        .burst_size(5)
        .key_extractor(SmartIpKeyExtractor)
        .finish()
        .expect(
            "failed to create governor config for sign in. this should not happen and is a bug",
        );

    let create_account_limiter_time: Option<String> =
        env::var("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND").ok();
    let create_account_limiter_burst: Option<String> =
        env::var("GATEKEEPER_CREATE_ACCOUNT_BURST").ok();

    //Default should be 608 requests per 5 minutes, PDS is 300 per 500 so will never hit it ideally
    let mut create_account_governor_conf = GovernorConfigBuilder::default();
    if create_account_limiter_time.is_some() {
        let time = create_account_limiter_time
            .expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND not set")
            .parse::<u64>()
            .expect("GATEKEEPER_CREATE_ACCOUNT_PER_SECOND must be a valid integer");
        create_account_governor_conf.per_second(time);
    }

    if create_account_limiter_burst.is_some() {
        let burst = create_account_limiter_burst
            .expect("GATEKEEPER_CREATE_ACCOUNT_BURST not set")
            .parse::<u32>()
            .expect("GATEKEEPER_CREATE_ACCOUNT_BURST must be a valid integer");
        create_account_governor_conf.burst_size(burst);
    }

    let create_account_governor_conf = create_account_governor_conf
        .key_extractor(SmartIpKeyExtractor)
        .finish().expect(
        "failed to create governor config for create account. this should not happen and is a bug",
    );

    let create_session_governor_limiter = create_session_governor_conf.limiter().clone();
    let sign_in_governor_limiter = sign_in_governor_conf.limiter().clone();
    let create_account_governor_limiter = create_account_governor_conf.limiter().clone();

    let interval = Duration::from_secs(60);
    // a separate background task to clean up
    std::thread::spawn(move || {
        loop {
            std::thread::sleep(interval);
            create_session_governor_limiter.retain_recent();
            sign_in_governor_limiter.retain_recent();
            create_account_governor_limiter.retain_recent();
        }
    });

    let cors = CorsLayer::new()
        .allow_origin(Any)
        .allow_methods([Method::GET, Method::OPTIONS, Method::POST])
        .allow_headers(Any);

    let app = Router::new()
        .route("/", get(root_handler))
        .route("/xrpc/com.atproto.server.getSession", get(get_session))
        .route(
            "/xrpc/com.atproto.server.updateEmail",
            post(update_email).layer(ax_middleware::from_fn(middleware::extract_did)),
        )
        .route(
            "/@atproto/oauth-provider/~api/sign-in",
            post(sign_in).layer(GovernorLayer::new(sign_in_governor_conf)),
        )
        .route(
            "/xrpc/com.atproto.server.createSession",
            post(create_session.layer(GovernorLayer::new(create_session_governor_conf))),
        )
        .route(
            "/xrpc/com.atproto.server.createAccount",
            post(create_account).layer(GovernorLayer::new(create_account_governor_conf)),
        )
        .layer(CompressionLayer::new())
        .layer(cors)
        .with_state(state);

    let host = env::var("GATEKEEPER_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let port: u16 = env::var("GATEKEEPER_PORT")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(8080);
    let addr: SocketAddr = format!("{host}:{port}")
        .parse()
        .expect("valid socket address");

    let listener = tokio::net::TcpListener::bind(addr).await?;

    let server = axum::serve(
        listener,
        app.into_make_service_with_connect_info::<SocketAddr>(),
    )
    .with_graceful_shutdown(shutdown_signal());

    if let Err(err) = server.await {
        log::error!("server error:{err}");
    }

    Ok(())
}

fn setup_tracing() {
    let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
    tracing_subscriber::registry()
        .with(env_filter)
        .with(fmt::layer())
        .init();
}

async fn shutdown_signal() {
    // Wait for Ctrl+C
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        use tokio::signal::unix::{SignalKind, signal};

        let mut sigterm =
            signal(SignalKind::terminate()).expect("failed to install signal handler");
        sigterm.recv().await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {},
        _ = terminate => {},
    }
}