use std::fmt;
use std::future::pending;
use std::io;
use std::sync::Arc;
use std::time::{Duration, SystemTime, SystemTimeError, UNIX_EPOCH};
use axum::body::{Body, Bytes, to_bytes};
use axum::extract::{DefaultBodyLimit, Path, Request, State};
use axum::http::header::AUTHORIZATION;
use axum::http::{HeaderMap, Method, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::{delete, get, post};
use axum::serve as serve_router;
use axum::{Json, Router};
use reqwest::redirect::Policy;
use reqwest::{Client, Error as ReqwestError};
use serde_json::{from_slice, json};
use thiserror::Error;
use tokio::net::TcpListener;
use tokio::signal::ctrl_c;
#[cfg(unix)]
use tokio::signal::unix::{SignalKind, signal};
use tokio::task::JoinHandle;
use tokio::{select, spawn, time};
use tracing::{info, warn};
use url::Url;
use ironflow_core::auth_proxy::{
AuthProxyError, AuthProxyRegistry, DEFAULT_UPSTREAM, TokenRejection, TokenRequest,
admin_key_matches, downstream_headers, error_body, extract_opaque_token, is_allowed_method,
is_allowed_path, upstream_headers,
};
use ironflow_store::crypto::{CryptoError, KeyRing};
use ironflow_store::error::StoreError;
use ironflow_store::postgres::PostgresStore;
pub const DATABASE_URL_ENV: &str = "IRONFLOW_AUTH_PROXY_DATABASE_URL";
pub const MIN_ADMIN_KEY_LEN: usize = 32;
pub const DEFAULT_MAX_BODY_BYTES: usize = 32 * 1024 * 1024;
const SHORT_ID_LEN: usize = 12;
const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(Clone)]
pub struct AuthProxyConfig {
pub upstream: Url,
pub admin_key: String,
pub max_body_bytes: usize,
}
impl fmt::Debug for AuthProxyConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AuthProxyConfig")
.field("upstream", &self.upstream.as_str())
.field("admin_key", &"<redacted>")
.field("max_body_bytes", &self.max_body_bytes)
.finish()
}
}
impl AuthProxyConfig {
pub fn new(admin_key: &str) -> Self {
assert!(
admin_key.len() >= MIN_ADMIN_KEY_LEN,
"the auth proxy admin key must be at least {MIN_ADMIN_KEY_LEN} characters"
);
Self {
upstream: Url::parse(DEFAULT_UPSTREAM).expect("DEFAULT_UPSTREAM is a valid URL"),
admin_key: admin_key.to_string(),
max_body_bytes: DEFAULT_MAX_BODY_BYTES,
}
}
pub fn with_upstream(mut self, upstream: Url) -> Self {
self.upstream = upstream;
self
}
}
#[derive(Clone)]
pub struct AuthProxyState {
registry: AuthProxyRegistry,
config: Arc<AuthProxyConfig>,
http: Client,
}
impl AuthProxyState {
pub fn new(config: AuthProxyConfig) -> Result<Self, ReqwestError> {
Self::with_registry(config, AuthProxyRegistry::default())
}
pub fn with_registry(
config: AuthProxyConfig,
registry: AuthProxyRegistry,
) -> Result<Self, ReqwestError> {
let http = Client::builder()
.connect_timeout(UPSTREAM_CONNECT_TIMEOUT)
.redirect(Policy::none())
.build()?;
Ok(Self {
registry,
config: Arc::new(config),
http,
})
}
pub fn registry(&self) -> &AuthProxyRegistry {
&self.registry
}
}
#[derive(Debug, Error)]
pub enum RegistryConfigError {
#[error(
"{DATABASE_URL_ENV} is set but no encryption key is configured: set IRONFLOW_SECRET_KEYS (or IRONFLOW_SECRET_KEY)"
)]
MissingKeyRing,
#[error("invalid encryption key: {0}")]
Crypto(#[from] CryptoError),
#[error("cannot open the token registry database: {0}")]
Store(#[from] StoreError),
}
pub async fn registry_from_config(
database_url: Option<&str>,
key_ring: Option<KeyRing>,
) -> Result<AuthProxyRegistry, RegistryConfigError> {
let Some(url) = database_url.map(str::trim).filter(|url| !url.is_empty()) else {
return Ok(AuthProxyRegistry::default());
};
let ring = key_ring.ok_or(RegistryConfigError::MissingKeyRing)?;
let mut store = PostgresStore::new(url).await?;
store.set_key_ring(ring);
Ok(AuthProxyRegistry::with_backend(Arc::new(store)))
}
pub fn router(state: AuthProxyState) -> Router {
let limit = state.config.max_body_bytes;
Router::new()
.route("/healthz", get(healthz))
.route("/admin/v1/tokens", post(issue_token))
.route("/admin/v1/tokens/{id}", delete(revoke_token))
.route("/admin/v1/runs/{run_id}/tokens", delete(revoke_run))
.fallback(relay)
.layer(DefaultBodyLimit::max(limit))
.with_state(state)
}
pub async fn serve(listener: TcpListener, state: AuthProxyState) -> io::Result<()> {
serve_router(listener, router(state))
.with_graceful_shutdown(shutdown_signal())
.await
}
pub fn spawn_purge(registry: AuthProxyRegistry, interval: Duration) -> JoinHandle<()> {
assert!(
!interval.is_zero(),
"purge interval must be greater than zero"
);
spawn(async move {
let mut ticker = time::interval(interval);
loop {
ticker.tick().await;
match now_unix() {
Ok(now) => match registry.purge_expired(now).await {
Ok(purged) if purged > 0 => info!(purged, "expired auth proxy tokens purged"),
Ok(_) => {}
Err(e) => warn!(error = %e, "expired auth proxy tokens purge failed"),
},
Err(e) => warn!(error = %e, "clock before the unix epoch; purge skipped"),
}
}
})
}
async fn shutdown_signal() {
let interrupt = async {
if let Err(e) = ctrl_c().await {
warn!(error = %e, "cannot listen for ctrl-c");
pending::<()>().await;
}
};
#[cfg(unix)]
let terminate = async {
match signal(SignalKind::terminate()) {
Ok(mut sigterm) => {
sigterm.recv().await;
}
Err(e) => {
warn!(error = %e, "cannot listen for SIGTERM");
pending::<()>().await;
}
}
};
#[cfg(not(unix))]
let terminate = pending::<()>();
select! {
() = interrupt => {},
() = terminate => {},
}
info!("shutting down");
}
fn now_unix() -> Result<u64, SystemTimeError> {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
}
fn short_id(id: &str) -> &str {
id.get(..SHORT_ID_LEN).unwrap_or(id)
}
fn error_response(status: StatusCode, kind: &str, message: &str) -> Response {
(status, Json(error_body(kind, message))).into_response()
}
fn clock_error(e: &SystemTimeError) -> Response {
warn!(error = %e, "system clock is before the unix epoch");
error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"api_error",
"auth proxy clock error",
)
}
fn admin_authorized(state: &AuthProxyState, headers: &HeaderMap) -> bool {
headers
.get(AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.strip_prefix("Bearer "))
.is_some_and(|key| admin_key_matches(&state.config.admin_key, key.trim()))
}
fn invalid_token(reason: &str, path: &str) -> Response {
warn!(reason, path = %path, "request with an invalid token rejected");
error_response(
StatusCode::UNAUTHORIZED,
"authentication_error",
"invalid or expired ironflow auth proxy token",
)
}
fn registry_unavailable(message: &str) -> Response {
error_response(StatusCode::SERVICE_UNAVAILABLE, "api_error", message)
}
fn admin_unauthorized() -> Response {
warn!("admin request without a valid admin key");
error_response(
StatusCode::UNAUTHORIZED,
"authentication_error",
"invalid or missing admin key",
)
}
async fn healthz() -> &'static str {
"ok"
}
async fn issue_token(
State(state): State<AuthProxyState>,
headers: HeaderMap,
body: Bytes,
) -> Response {
if !admin_authorized(&state, &headers) {
return admin_unauthorized();
}
let Ok(request) = from_slice::<TokenRequest>(&body) else {
warn!("token request body rejected");
return error_response(
StatusCode::BAD_REQUEST,
"invalid_request_error",
"invalid token request body",
);
};
let now = match now_unix() {
Ok(now) => now,
Err(e) => return clock_error(&e),
};
let run_id = request.run_id.clone();
let step = request.step.clone();
match state.registry.issue(request, now).await {
Ok(issued) => {
info!(
token = %issued.short_id(),
run_id = %run_id,
step = %step,
"token issued"
);
(StatusCode::CREATED, Json(issued)).into_response()
}
Err(AuthProxyError::InvalidRequest(message)) => {
warn!(run_id = %run_id, step = %step, reason = %message, "token request refused");
error_response(StatusCode::BAD_REQUEST, "invalid_request_error", &message)
}
Err(AuthProxyError::Backend(e)) => {
warn!(run_id = %run_id, step = %step, error = %e, "token registry unavailable");
registry_unavailable("token registry unavailable")
}
Err(e) => {
warn!(error = %e, "token issuance failed");
error_response(
StatusCode::INTERNAL_SERVER_ERROR,
"api_error",
"token issuance failed",
)
}
}
}
async fn revoke_token(
State(state): State<AuthProxyState>,
Path(id): Path<String>,
headers: HeaderMap,
) -> Response {
if !admin_authorized(&state, &headers) {
return admin_unauthorized();
}
match state.registry.revoke(&id).await {
Ok(true) => {
info!(token = %short_id(&id), "token revoked");
StatusCode::NO_CONTENT.into_response()
}
Ok(false) => error_response(StatusCode::NOT_FOUND, "not_found_error", "unknown token"),
Err(e) => {
warn!(token = %short_id(&id), error = %e, "token registry unavailable");
registry_unavailable("token registry unavailable")
}
}
}
async fn revoke_run(
State(state): State<AuthProxyState>,
Path(run_id): Path<String>,
headers: HeaderMap,
) -> Response {
if !admin_authorized(&state, &headers) {
return admin_unauthorized();
}
match state.registry.revoke_run(&run_id).await {
Ok(revoked) => {
info!(run_id = %run_id, revoked, "run tokens revoked");
(StatusCode::OK, Json(json!({ "revoked": revoked }))).into_response()
}
Err(e) => {
warn!(run_id = %run_id, error = %e, "token registry unavailable");
registry_unavailable("token registry unavailable")
}
}
}
async fn relay(State(state): State<AuthProxyState>, request: Request) -> Response {
let (parts, body) = request.into_parts();
let method = parts.method;
let path = parts.uri.path().to_string();
if parts.uri.authority().is_some() || method == Method::CONNECT {
warn!(method = %method, "request for another host refused");
return error_response(
StatusCode::FORBIDDEN,
"permission_error",
"only api.anthropic.com is reachable through this proxy",
);
}
let now = match now_unix() {
Ok(now) => now,
Err(e) => return clock_error(&e),
};
let Some(token) = extract_opaque_token(&parts.headers) else {
return invalid_token("missing", &path);
};
let grant = match state.registry.resolve(&token, now).await {
Ok(grant) => grant,
Err(TokenRejection::Unknown) => return invalid_token("unknown", &path),
Err(TokenRejection::Expired) => return invalid_token("expired", &path),
Err(TokenRejection::Unavailable(e)) => {
warn!(error = %e, path = %path, "token registry unavailable");
return registry_unavailable("auth proxy token registry unavailable");
}
};
let token = short_id(&grant.id);
if !is_allowed_path(&path) {
warn!(token = %token, path = %path, "path outside the API refused");
return error_response(
StatusCode::FORBIDDEN,
"permission_error",
"only the Anthropic API under /v1/ is reachable through this proxy",
);
}
if !is_allowed_method(&method) {
warn!(token = %token, method = %method, path = %path, "method refused");
return error_response(
StatusCode::METHOD_NOT_ALLOWED,
"invalid_request_error",
"only GET and POST are relayed",
);
}
let bytes = match to_bytes(body, state.config.max_body_bytes).await {
Ok(bytes) => bytes,
Err(e) => {
warn!(token = %token, error = %e, "request body rejected");
return error_response(
StatusCode::PAYLOAD_TOO_LARGE,
"request_too_large",
"request body too large or unreadable",
);
}
};
let mut url = state.config.upstream.clone();
url.set_path(&path);
url.set_query(parts.uri.query());
let upstream = state
.http
.request(method.clone(), url)
.headers(upstream_headers(&parts.headers, &grant.credential))
.body(bytes)
.send()
.await;
let upstream = match upstream {
Ok(upstream) => upstream,
Err(e) => {
warn!(token = %token, error = %e.without_url(), "upstream request failed");
return error_response(
StatusCode::BAD_GATEWAY,
"api_error",
"the Anthropic API is unreachable",
);
}
};
let status = upstream.status();
info!(
token = %token,
run_id = %grant.run_id,
step = %grant.step,
method = %method,
path = %path,
status = status.as_u16(),
"relayed"
);
let headers = downstream_headers(upstream.headers());
let mut response = Response::new(Body::from_stream(upstream.bytes_stream()));
*response.status_mut() = status;
*response.headers_mut() = headers;
response
}
#[cfg(test)]
mod tests {
use super::*;
fn key_ring() -> KeyRing {
KeyRing::from_spec(&format!("1:{}", "aa".repeat(32)), None).unwrap()
}
#[tokio::test]
async fn registry_from_config_defaults_to_memory() {
let registry = registry_from_config(None, None).await.unwrap();
assert!(registry.is_empty().await.unwrap());
let registry = registry_from_config(Some(""), None).await.unwrap();
assert!(registry.is_empty().await.unwrap());
let registry = registry_from_config(Some(" \n"), Some(key_ring()))
.await
.unwrap();
assert!(registry.is_empty().await.unwrap());
}
#[tokio::test]
async fn registry_from_config_requires_key_ring_with_database() {
let result = registry_from_config(Some("postgres://localhost/x"), None).await;
assert!(
matches!(result, Err(RegistryConfigError::MissingKeyRing)),
"{result:?}"
);
}
#[tokio::test]
async fn registry_from_config_rejects_invalid_database_url() {
let result = registry_from_config(Some("not a url"), Some(key_ring())).await;
assert!(
matches!(result, Err(RegistryConfigError::Store(_))),
"{result:?}"
);
}
}