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, CONTENT_TYPE};
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 metrics::counter;
use metrics_exporter_prometheus::PrometheusHandle;
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, Grant, GrantCredential, RELAY_PREFIX,
TokenRejection, TokenRequest, admin_key_matches, downstream_headers, error_body,
extract_opaque_token, is_allowed_method, is_allowed_path, is_relay_method, is_relay_path,
is_valid_request_host, secret_upstream_headers, 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;
pub const REQUESTS_TOTAL: &str = "ironflow_auth_proxy_requests_total";
const METRICS_CONTENT_TYPE: &str = "text/plain; version=0.0.4";
const SECRET_UPSTREAM_BASE: &str = "https://localhost";
const SHORT_ID_LEN: usize = 12;
const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
#[derive(Clone)]
pub struct AuthProxyConfig {
pub upstream: Url,
pub secret_upstream: Option<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(
"secret_upstream",
&self.secret_upstream.as_ref().map(Url::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"),
secret_upstream: None,
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
}
pub fn with_secret_upstream(mut self, upstream: Url) -> Self {
self.secret_upstream = Some(upstream);
self
}
}
#[derive(Clone)]
pub struct AuthProxyState {
registry: AuthProxyRegistry,
config: Arc<AuthProxyConfig>,
http: Client,
metrics: Option<PrometheusHandle>,
}
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,
metrics: None,
})
}
pub fn registry(&self) -> &AuthProxyRegistry {
&self.registry
}
pub fn with_metrics(mut self, handle: PrometheusHandle) -> Self {
self.metrics = Some(handle);
self
}
}
#[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("/metrics", get(serve_metrics))
.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");
unauthorized()
}
fn unauthorized() -> Response {
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 serve_metrics(State(state): State<AuthProxyState>) -> Response {
match &state.metrics {
Some(handle) => ([(CONTENT_TYPE, METRICS_CONTENT_TYPE)], handle.render()).into_response(),
None => error_response(
StatusCode::NOT_FOUND,
"not_found_error",
"metrics are disabled",
),
}
}
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();
let secret = match &request.credential {
GrantCredential::Secret(secret) => Some(secret.name().to_string()),
GrantCredential::Claude(_) => None,
};
match state.registry.issue(request, now).await {
Ok(issued) => {
info!(
token = %issued.short_id(),
run_id = %run_id,
step = %step,
secret = secret.as_deref(),
"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",
);
}
if let Some(rest) = path
.strip_prefix(RELAY_PREFIX)
.and_then(|rest| rest.strip_prefix('/'))
{
let query = parts.uri.query();
return relay_secret(&state, method, &parts.headers, query, body, rest).await;
}
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::Revoked) => return invalid_token("revoked", &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);
let GrantCredential::Claude(credential) = &grant.credential else {
warn!(token = %token, path = %path, "secret token on the Anthropic API refused");
return error_response(
StatusCode::FORBIDDEN,
"permission_error",
"this token only reaches its allowlisted hosts under /r/",
);
};
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, 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
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RelayResult {
Relayed,
ForbiddenHost,
ForbiddenPath,
ForbiddenMethod,
UnknownToken,
Expired,
Revoked,
UpstreamError,
Unavailable,
}
impl RelayResult {
fn as_str(self) -> &'static str {
match self {
Self::Relayed => "relayed",
Self::ForbiddenHost => "forbidden_host",
Self::ForbiddenPath => "forbidden_path",
Self::ForbiddenMethod => "forbidden_method",
Self::UnknownToken => "unknown_token",
Self::Expired => "expired",
Self::Revoked => "revoked",
Self::UpstreamError => "upstream_error",
Self::Unavailable => "unavailable",
}
}
}
fn record(
result: RelayResult,
host: &str,
secret: &str,
grant: Option<&Grant>,
upstream_status: Option<u16>,
) {
counter!(REQUESTS_TOTAL, "secret" => secret.to_string(), "result" => result.as_str())
.increment(1);
let token = grant.map(|grant| short_id(&grant.id));
let run_id = grant.map(|grant| grant.run_id.as_str());
let step = grant.map(|grant| grant.step.as_str());
if result == RelayResult::Relayed {
info!(
host,
secret,
run_id,
step,
token,
result = result.as_str(),
upstream_status,
"secret relay"
);
} else {
warn!(
host,
secret,
run_id,
step,
token,
result = result.as_str(),
upstream_status,
"secret relay"
);
}
}
async fn relay_secret(
state: &AuthProxyState,
method: Method,
headers: &HeaderMap,
query: Option<&str>,
body: Body,
rest: &str,
) -> Response {
let (host, path) = match rest.split_once('/') {
Some((host, path)) => (host.to_ascii_lowercase(), format!("/{path}")),
None => (rest.to_ascii_lowercase(), "/".to_string()),
};
let logged_host = if is_valid_request_host(&host) {
host.as_str()
} else {
"<invalid>"
};
let now = match now_unix() {
Ok(now) => now,
Err(e) => return clock_error(&e),
};
let Some(token) = extract_opaque_token(headers) else {
record(RelayResult::UnknownToken, logged_host, "", None, None);
return unauthorized();
};
let grant = match state.registry.resolve(&token, now).await {
Ok(grant) => grant,
Err(rejection) => {
let result = match rejection {
TokenRejection::Unknown => RelayResult::UnknownToken,
TokenRejection::Expired => RelayResult::Expired,
TokenRejection::Revoked => RelayResult::Revoked,
TokenRejection::Unavailable(e) => {
warn!(error = %e, "token registry unavailable");
record(RelayResult::Unavailable, logged_host, "", None, None);
return registry_unavailable("auth proxy token registry unavailable");
}
};
record(result, logged_host, "", None, None);
return unauthorized();
}
};
let GrantCredential::Secret(secret) = &grant.credential else {
record(
RelayResult::ForbiddenHost,
logged_host,
"",
Some(&grant),
None,
);
return error_response(
StatusCode::FORBIDDEN,
"permission_error",
"a Claude token only reaches the Anthropic API",
);
};
let name = secret.name();
let refuse = |result| record(result, logged_host, name, Some(&grant), None);
let forbidden_host = || {
refuse(RelayResult::ForbiddenHost);
error_response(
StatusCode::FORBIDDEN,
"permission_error",
"host not allowed for this secret",
)
};
if !secret.allows_host(&host) {
return forbidden_host();
}
if !is_relay_path(&path) {
refuse(RelayResult::ForbiddenPath);
return error_response(
StatusCode::FORBIDDEN,
"permission_error",
"path not allowed: no '..', '//' or percent-encoding",
);
}
if !is_relay_method(&method) {
refuse(RelayResult::ForbiddenMethod);
return error_response(
StatusCode::METHOD_NOT_ALLOWED,
"invalid_request_error",
"only GET, HEAD, POST, PUT, PATCH and DELETE are relayed",
);
}
let bytes = match to_bytes(body, state.config.max_body_bytes).await {
Ok(bytes) => bytes,
Err(e) => {
warn!(
token = %short_id(&grant.id),
secret = name,
error = %e,
"request body rejected"
);
return error_response(
StatusCode::PAYLOAD_TOO_LARGE,
"request_too_large",
"request body too large or unreadable",
);
}
};
let mut url = match &state.config.secret_upstream {
Some(upstream) => upstream.clone(),
None => {
let url = Url::parse(SECRET_UPSTREAM_BASE).ok().and_then(|mut url| {
url.set_host(Some(&host)).ok()?;
Some(url)
});
let Some(url) = url else {
return forbidden_host();
};
url
}
};
url.set_path(&path);
url.set_query(query);
let upstream = state
.http
.request(method, url)
.headers(secret_upstream_headers(headers, secret))
.body(bytes)
.send()
.await;
let upstream = match upstream {
Ok(upstream) => upstream,
Err(e) => {
warn!(
token = %short_id(&grant.id),
secret = name,
error = %e.without_url(),
"upstream request failed"
);
refuse(RelayResult::UpstreamError);
return error_response(
StatusCode::BAD_GATEWAY,
"upstream_error",
"upstream unreachable",
);
}
};
let status = upstream.status();
record(
RelayResult::Relayed,
logged_host,
name,
Some(&grant),
Some(status.as_u16()),
);
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:?}"
);
}
#[test]
fn relay_results_have_stable_labels() {
let labels: Vec<&str> = [
RelayResult::Relayed,
RelayResult::ForbiddenHost,
RelayResult::ForbiddenPath,
RelayResult::ForbiddenMethod,
RelayResult::UnknownToken,
RelayResult::Expired,
RelayResult::Revoked,
RelayResult::UpstreamError,
RelayResult::Unavailable,
]
.into_iter()
.map(RelayResult::as_str)
.collect();
assert_eq!(
labels,
vec![
"relayed",
"forbidden_host",
"forbidden_path",
"forbidden_method",
"unknown_token",
"expired",
"revoked",
"upstream_error",
"unavailable",
]
);
}
#[test]
fn config_debug_shows_secret_upstream_not_admin_key() {
let config = AuthProxyConfig::new("0123456789abcdef0123456789abcdef")
.with_secret_upstream(Url::parse("http://127.0.0.1:9001").unwrap());
let debug = format!("{config:?}");
assert!(debug.contains("http://127.0.0.1:9001/"), "{debug}");
assert!(!debug.contains("0123456789abcdef"), "{debug}");
}
}