Something went wrong. Try again.
An OIDC provider for AT Protocol OAuth. Rust port of apenwarr/atlogin. forked from quilling.dev/ratlogin
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291use std::net::{IpAddr, SocketAddr};use std::sync::Arc;use std::time::Duration;
use axum::Router;use axum::extract::ConnectInfo;use axum::http::Request;use jiff::Timestamp;use tokio::net::TcpListener;use tokio::signal::unix::{SignalKind, signal};use tower_governor::GovernorLayer;use tower_governor::errors::GovernorError;use tower_governor::governor::GovernorConfigBuilder;use tower_governor::key_extractor::KeyExtractor;use tower_http::limit::RequestBodyLimitLayer;use tower_http::timeout::TimeoutLayer;
use crate::atproto::{AtprotoLogin, JacquardLogin, LoginError};use crate::authstore::SqliteAuthStore;use crate::clients::ClientStore;use crate::config::{RateLimit, ServerConfig};use crate::database::{self, Database, DatabaseError};use crate::state::StateKeys;use crate::web::{AppState, AppStateError, router};
pub const SWEEP_INTERVAL: Duration = Duration::from_secs(60);pub const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);pub const BODY_LIMIT_BYTES: usize = 64 * 1024;
#[derive(Debug, thiserror::Error)]pub enum ServeError { #[error(transparent)] Login(#[from] LoginError), #[error("cannot listen on {listen}: {source}")] Bind { listen: String, source: std::io::Error, }, #[error("server failed: {0}")] Io(#[from] std::io::Error), #[error(transparent)] Database(#[from] DatabaseError), #[error(transparent)] State(#[from] AppStateError), #[error("invalid rate limit configuration")] RateLimit,}
pub fn init_tracing() { let filter = tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); let _ = tracing_subscriber::fmt() .with_env_filter(filter) .with_writer(std::io::stderr) .try_init();}
pub async fn serve(config: ServerConfig, keys: StateKeys) -> Result<(), ServeError> { let db = Database::open(&config.state_dir.join(database::FILE_NAME))?; let login = JacquardLogin::new( &config, SqliteAuthStore::new(db.clone()), keys.client.keyset(), )?; let shutdown = shutdown_signal()?; let listener = TcpListener::bind(&config.listen) .await .map_err(|source| ServeError::Bind { listen: config.listen.clone(), source, })?; run(listener, config, keys, login, db, shutdown).await}
pub async fn run<L: AtprotoLogin + 'static>( listener: TcpListener, config: ServerConfig, keys: StateKeys, login: L, db: Database, shutdown: impl Future<Output = ()> + Send + 'static,) -> Result<(), ServeError> { warn_about_plaintext_redirects(&config, &db); let listen = listener.local_addr()?; let issuer = config.issuer().to_owned(); let trust_proxy_headers = config.trust_proxy_headers; let rate_limit = config.rate_limit; let state = Arc::new(AppState::new( config, keys.signing, &keys.client, login, db, )?);
let sweeper = Arc::clone(&state); tokio::spawn(async move { let mut interval = tokio::time::interval(SWEEP_INTERVAL); loop { interval.tick().await; if let Err(error) = sweeper.flows.sweep(Timestamp::now()) { tracing::error!(%error, "sweep failed"); } } });
let app = limited(router(state), trust_proxy_headers, rate_limit)?; tracing::info!( version = env!("CARGO_PKG_VERSION"), %listen, %issuer, %trust_proxy_headers, "listening" ); if !trust_proxy_headers && listen.ip().is_loopback() { tracing::warn!( "listening on loopback without trust_proxy_headers, every request shares the proxy's one rate-limit key" ); } axum::serve( listener, app.into_make_service_with_connect_info::<SocketAddr>(), ) .with_graceful_shutdown(shutdown) .await?; tracing::info!("stopped"); Ok(())}
fn warn_about_plaintext_redirects(config: &ServerConfig, db: &Database) { let stored = ClientStore::new(db.clone()) .active() .unwrap_or_else(|error| { tracing::error!(%error, "cannot list stored clients"); Vec::new() }); for client in config.clients.clients().chain(&stored) { let uris = client.plaintext_redirect_uris(); if !uris.is_empty() { tracing::warn!( client = client.id().as_str(), redirect_uris = uris.join(" "), "client accepts redirect URIs that are not HTTPS" ); } }}
fn limited( app: Router, trust_proxy_headers: bool, rate_limit: RateLimit,) -> Result<Router, ServeError> { let app = app .layer(TimeoutLayer::with_status_code( http::StatusCode::REQUEST_TIMEOUT, REQUEST_TIMEOUT, )) .layer(RequestBodyLimitLayer::new(BODY_LIMIT_BYTES)); let config = GovernorConfigBuilder::default() .per_second(rate_limit.period_seconds) .burst_size(rate_limit.burst) .key_extractor(RateLimitKey { trust_proxy_headers, }) .finish() .ok_or(ServeError::RateLimit)?; let limiter = config.limiter().clone(); tokio::spawn(async move { let mut interval = tokio::time::interval(SWEEP_INTERVAL); loop { interval.tick().await; limiter.retain_recent(); } }); Ok(app.layer(GovernorLayer::new(config)))}
fn shutdown_signal() -> std::io::Result<impl Future<Output = ()> + Send + 'static> { let mut terminate = signal(SignalKind::terminate())?; let mut interrupt = signal(SignalKind::interrupt())?; Ok(async move { tokio::select! { _ = terminate.recv() => {} _ = interrupt.recv() => {} } })}
#[derive(Clone)]struct RateLimitKey { trust_proxy_headers: bool,}
impl KeyExtractor for RateLimitKey { type Key = IpAddr;
fn extract<T>(&self, request: &Request<T>) -> Result<IpAddr, GovernorError> { let forwarded = request .headers() .get_all("x-forwarded-for") .iter() .filter_map(|value| value.to_str().ok()) .flat_map(|value| value.split(',')) .filter_map(|entry| entry.trim().parse::<IpAddr>().ok()) .next_back() .filter(|_| self.trust_proxy_headers); forwarded .or_else(|| { request .extensions() .get::<ConnectInfo<SocketAddr>>() .map(|peer| peer.ip()) }) .map(network) .ok_or(GovernorError::UnableToExtractKey) }}
fn network(address: IpAddr) -> IpAddr { match address { IpAddr::V4(_) => address, IpAddr::V6(v6) => IpAddr::V6((u128::from(v6) & !((1u128 << 64) - 1)).into()), }}
#[cfg(test)]mod tests { use super::*;
fn keyed(forwarded: &[&str], peer: &str) -> IpAddr { keyed_by(true, forwarded, peer) }
fn keyed_by(trust_proxy_headers: bool, forwarded: &[&str], peer: &str) -> IpAddr { let mut request = Request::new(()); for value in forwarded { request .headers_mut() .append("x-forwarded-for", value.parse().expect("header")); } request .extensions_mut() .insert(ConnectInfo(peer.parse::<SocketAddr>().expect("peer"))); RateLimitKey { trust_proxy_headers, } .extract(&request) .expect("key") }
#[test] fn the_rate_limit_key_is_the_address_the_proxy_appended() { assert_eq!( keyed(&[], "10.0.0.1:4000"), "10.0.0.1".parse::<IpAddr>().unwrap() ); assert_eq!( keyed(&["203.0.113.9, 198.51.100.7"], "10.0.0.1:4000"), "198.51.100.7".parse::<IpAddr>().unwrap() ); assert_eq!( keyed(&["203.0.113.9", "198.51.100.7"], "10.0.0.1:4000"), "198.51.100.7".parse::<IpAddr>().unwrap() ); assert_eq!( keyed(&["garbage"], "10.0.0.1:4000"), "10.0.0.1".parse::<IpAddr>().unwrap() ); assert_eq!( keyed_by(false, &["203.0.113.9"], "10.0.0.1:4000"), "10.0.0.1".parse::<IpAddr>().unwrap(), "forwarded headers are ignored without a trusted proxy" ); }
#[test] fn ipv6_visitors_share_one_key_per_64() { let first = keyed(&[], "[2001:db8:1:2:aaaa::1]:4000"); let second = keyed(&[], "[2001:db8:1:2:ffff:ffff:ffff:ffff]:4000"); let other = keyed(&[], "[2001:db8:1:3::1]:4000"); assert_eq!(first, "2001:db8:1:2::".parse::<IpAddr>().unwrap()); assert_eq!(first, second); assert_ne!(first, other); assert_eq!( keyed(&["2001:db8:9:9::5"], "10.0.0.1:4000"), "2001:db8:9:9::".parse::<IpAddr>().unwrap() ); }}