#![warn(clippy::unwrap_used)] use crate::oauth_provider::sign_in; use crate::xrpc::com_atproto_server::{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_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; #[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, mailer_from: String, template_engine: Engine>, } 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> { setup_tracing(); //TODO may need to change where this reads from? Like an env variable for it's location? Or arg? dotenvy::from_path(Path::new("./pds.env"))?; let pds_root = env::var("PDS_DATA_DIRECTORY")?; 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 = AsyncSmtpTransport::::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::(); } 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) .finish() .expect("failed to create governor config. 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) .finish() .expect("failed to create governor config. 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 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(); } }); 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).layer(ax_middleware::from_fn(middleware::extract_did)), ) .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))), ) .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::(), ) .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 => {}, } }