ironflow_auth_proxy/
lib.rs

1//! # ironflow-auth-proxy
2//!
3//! HTTP service that keeps the Claude credential out of ironflow agent pods.
4//!
5//! The worker (`K8sEphemeralProvider::auth_proxy`) asks the admin API for an
6//! opaque token bound to one run and one step, and hands only that token to
7//! the pod. Claude Code sends it as `Authorization: Bearer` to this proxy,
8//! which swaps it for the real credential and relays the request to
9//! `api.anthropic.com`, streaming the answer back.
10//!
11//! One listener serves:
12//!
13//! * `GET /healthz` - liveness;
14//! * `/admin/v1/...` - token issuance and revocation, behind
15//!   `Authorization: Bearer <IRONFLOW_AUTH_PROXY_ADMIN_KEY>`;
16//! * anything else - the relay: an unknown, expired or revoked token gets a
17//!   401, a path outside `/v1/` or a request for another host a 403, a method
18//!   other than GET/POST a 405.
19//!
20//! Tokens live in memory: run a single replica. Logs never contain a token or
21//! a credential, only the short token id.
22//!
23//! # Examples
24//!
25//! ```no_run
26//! use std::time::Duration;
27//!
28//! use ironflow_auth_proxy::{AuthProxyConfig, AuthProxyState, serve, spawn_purge};
29//! use tokio::net::TcpListener;
30//!
31//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
32//! let state = AuthProxyState::new(AuthProxyConfig::new("0123456789abcdef0123456789abcdef"))?;
33//! let purge = spawn_purge(state.registry().clone(), Duration::from_secs(60));
34//! let listener = TcpListener::bind("0.0.0.0:8080").await?;
35//! serve(listener, state).await?;
36//! purge.abort();
37//! # Ok(())
38//! # }
39//! ```
40
41use std::fmt;
42use std::future::pending;
43use std::io;
44use std::sync::Arc;
45use std::time::{Duration, SystemTime, SystemTimeError, UNIX_EPOCH};
46
47use axum::body::{Body, Bytes, to_bytes};
48use axum::extract::{DefaultBodyLimit, Path, Request, State};
49use axum::http::header::AUTHORIZATION;
50use axum::http::{HeaderMap, Method, StatusCode};
51use axum::response::{IntoResponse, Response};
52use axum::routing::{delete, get, post};
53use axum::serve as serve_router;
54use axum::{Json, Router};
55use reqwest::redirect::Policy;
56use reqwest::{Client, Error as ReqwestError};
57use serde_json::{from_slice, json};
58use tokio::net::TcpListener;
59use tokio::signal::ctrl_c;
60#[cfg(unix)]
61use tokio::signal::unix::{SignalKind, signal};
62use tokio::task::JoinHandle;
63use tokio::{select, spawn, time};
64use tracing::{info, warn};
65use url::Url;
66
67use ironflow_core::auth_proxy::{
68    AuthProxyError, AuthProxyRegistry, DEFAULT_UPSTREAM, TokenRejection, TokenRequest,
69    admin_key_matches, downstream_headers, error_body, extract_opaque_token, is_allowed_method,
70    is_allowed_path, upstream_headers,
71};
72
73/// Shortest admin key accepted.
74pub const MIN_ADMIN_KEY_LEN: usize = 32;
75
76/// Default largest request body relayed: 32 MiB.
77pub const DEFAULT_MAX_BODY_BYTES: usize = 32 * 1024 * 1024;
78
79/// Length of the token id prefix written to logs.
80const SHORT_ID_LEN: usize = 12;
81
82/// Connect timeout towards the upstream. There is no total timeout: an
83/// agent turn streams for as long as the model answers.
84const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
85
86/// Configuration of the proxy. Its [`Debug`] output never shows the admin key.
87///
88/// # Examples
89///
90/// ```
91/// use ironflow_auth_proxy::AuthProxyConfig;
92///
93/// let config = AuthProxyConfig::new("0123456789abcdef0123456789abcdef");
94/// assert_eq!(config.upstream.as_str(), "https://api.anthropic.com/");
95/// assert!(!format!("{config:?}").contains("0123456789abcdef"));
96/// ```
97#[derive(Clone)]
98pub struct AuthProxyConfig {
99    /// Where requests are relayed. Always `https://api.anthropic.com` in the binary.
100    pub upstream: Url,
101    /// Key protecting the admin API.
102    pub admin_key: String,
103    /// Largest request body relayed, in bytes.
104    pub max_body_bytes: usize,
105}
106
107impl fmt::Debug for AuthProxyConfig {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        f.debug_struct("AuthProxyConfig")
110            .field("upstream", &self.upstream.as_str())
111            .field("admin_key", &"<redacted>")
112            .field("max_body_bytes", &self.max_body_bytes)
113            .finish()
114    }
115}
116
117impl AuthProxyConfig {
118    /// Configuration relaying to [`DEFAULT_UPSTREAM`] with the given admin key.
119    ///
120    /// # Panics
121    ///
122    /// Panics if `admin_key` is shorter than [`MIN_ADMIN_KEY_LEN`].
123    ///
124    /// # Examples
125    ///
126    /// See [`AuthProxyConfig`].
127    pub fn new(admin_key: &str) -> Self {
128        assert!(
129            admin_key.len() >= MIN_ADMIN_KEY_LEN,
130            "the auth proxy admin key must be at least {MIN_ADMIN_KEY_LEN} characters"
131        );
132        Self {
133            upstream: Url::parse(DEFAULT_UPSTREAM).expect("DEFAULT_UPSTREAM is a valid URL"),
134            admin_key: admin_key.to_string(),
135            max_body_bytes: DEFAULT_MAX_BODY_BYTES,
136        }
137    }
138
139    /// Relay to another upstream. **For tests only**: the binary never calls
140    /// it, so a deployed proxy only reaches `api.anthropic.com`.
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// use ironflow_auth_proxy::AuthProxyConfig;
146    /// use url::Url;
147    ///
148    /// # fn example() -> Result<(), url::ParseError> {
149    /// let config = AuthProxyConfig::new("0123456789abcdef0123456789abcdef")
150    ///     .with_upstream(Url::parse("http://127.0.0.1:9000")?);
151    /// assert_eq!(config.upstream.port(), Some(9000));
152    /// # Ok(())
153    /// # }
154    /// ```
155    pub fn with_upstream(mut self, upstream: Url) -> Self {
156        self.upstream = upstream;
157        self
158    }
159}
160
161/// Shared state of the proxy: the token registry, the configuration and the
162/// upstream HTTP client. Cheap to clone.
163///
164/// # Examples
165///
166/// ```
167/// use ironflow_auth_proxy::{AuthProxyConfig, AuthProxyState};
168///
169/// # fn example() -> Result<(), reqwest::Error> {
170/// let state = AuthProxyState::new(AuthProxyConfig::new("0123456789abcdef0123456789abcdef"))?;
171/// assert!(state.registry().is_empty());
172/// # Ok(())
173/// # }
174/// ```
175#[derive(Clone)]
176pub struct AuthProxyState {
177    registry: AuthProxyRegistry,
178    config: Arc<AuthProxyConfig>,
179    http: Client,
180}
181
182impl AuthProxyState {
183    /// Build the state with an empty registry.
184    ///
185    /// The upstream client never follows redirects and only bounds the
186    /// connection time.
187    ///
188    /// # Errors
189    ///
190    /// Returns the [`reqwest::Error`] raised when the HTTP client cannot be
191    /// built (TLS backend initialisation).
192    ///
193    /// # Examples
194    ///
195    /// See [`AuthProxyState`].
196    pub fn new(config: AuthProxyConfig) -> Result<Self, ReqwestError> {
197        let http = Client::builder()
198            .connect_timeout(UPSTREAM_CONNECT_TIMEOUT)
199            .redirect(Policy::none())
200            .build()?;
201        Ok(Self {
202            registry: AuthProxyRegistry::default(),
203            config: Arc::new(config),
204            http,
205        })
206    }
207
208    /// The token registry.
209    ///
210    /// # Examples
211    ///
212    /// See [`AuthProxyState`].
213    pub fn registry(&self) -> &AuthProxyRegistry {
214        &self.registry
215    }
216}
217
218/// The proxy router: health, admin API and relay.
219///
220/// # Examples
221///
222/// ```no_run
223/// use axum::serve;
224/// use ironflow_auth_proxy::{AuthProxyConfig, AuthProxyState, router};
225/// use tokio::net::TcpListener;
226///
227/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
228/// let state = AuthProxyState::new(AuthProxyConfig::new("0123456789abcdef0123456789abcdef"))?;
229/// let listener = TcpListener::bind("127.0.0.1:0").await?;
230/// serve(listener, router(state)).await?;
231/// # Ok(())
232/// # }
233/// ```
234pub fn router(state: AuthProxyState) -> Router {
235    let limit = state.config.max_body_bytes;
236    Router::new()
237        .route("/healthz", get(healthz))
238        .route("/admin/v1/tokens", post(issue_token))
239        .route("/admin/v1/tokens/{id}", delete(revoke_token))
240        .route("/admin/v1/runs/{run_id}/tokens", delete(revoke_run))
241        .fallback(relay)
242        .layer(DefaultBodyLimit::max(limit))
243        .with_state(state)
244}
245
246/// Serve the proxy on `listener` until ctrl-c or SIGTERM, letting in-flight
247/// requests finish.
248///
249/// # Errors
250///
251/// Returns the I/O error that stopped the server.
252///
253/// # Examples
254///
255/// See the [crate documentation](crate).
256pub async fn serve(listener: TcpListener, state: AuthProxyState) -> io::Result<()> {
257    serve_router(listener, router(state))
258        .with_graceful_shutdown(shutdown_signal())
259        .await
260}
261
262/// Spawn a task dropping the expired grants of `registry` every `interval`.
263///
264/// Abort the returned handle to stop it.
265///
266/// # Panics
267///
268/// Panics if `interval` is zero, or when called outside a Tokio runtime.
269///
270/// # Examples
271///
272/// ```no_run
273/// use std::time::Duration;
274///
275/// use ironflow_auth_proxy::spawn_purge;
276/// use ironflow_core::auth_proxy::AuthProxyRegistry;
277///
278/// # async fn example() {
279/// let purge = spawn_purge(AuthProxyRegistry::default(), Duration::from_secs(60));
280/// purge.abort();
281/// # }
282/// ```
283pub fn spawn_purge(registry: AuthProxyRegistry, interval: Duration) -> JoinHandle<()> {
284    assert!(
285        !interval.is_zero(),
286        "purge interval must be greater than zero"
287    );
288    spawn(async move {
289        let mut ticker = time::interval(interval);
290        loop {
291            ticker.tick().await;
292            match now_unix() {
293                Ok(now) => {
294                    let purged = registry.purge_expired(now);
295                    if purged > 0 {
296                        info!(purged, "expired auth proxy tokens purged");
297                    }
298                }
299                Err(e) => warn!(error = %e, "clock before the unix epoch; purge skipped"),
300            }
301        }
302    })
303}
304
305async fn shutdown_signal() {
306    let interrupt = async {
307        if let Err(e) = ctrl_c().await {
308            warn!(error = %e, "cannot listen for ctrl-c");
309            pending::<()>().await;
310        }
311    };
312    #[cfg(unix)]
313    let terminate = async {
314        match signal(SignalKind::terminate()) {
315            Ok(mut sigterm) => {
316                sigterm.recv().await;
317            }
318            Err(e) => {
319                warn!(error = %e, "cannot listen for SIGTERM");
320                pending::<()>().await;
321            }
322        }
323    };
324    #[cfg(not(unix))]
325    let terminate = pending::<()>();
326    select! {
327        () = interrupt => {},
328        () = terminate => {},
329    }
330    info!("shutting down");
331}
332
333fn now_unix() -> Result<u64, SystemTimeError> {
334    SystemTime::now()
335        .duration_since(UNIX_EPOCH)
336        .map(|d| d.as_secs())
337}
338
339fn short_id(id: &str) -> &str {
340    id.get(..SHORT_ID_LEN).unwrap_or(id)
341}
342
343fn error_response(status: StatusCode, kind: &str, message: &str) -> Response {
344    (status, Json(error_body(kind, message))).into_response()
345}
346
347fn clock_error(e: &SystemTimeError) -> Response {
348    warn!(error = %e, "system clock is before the unix epoch");
349    error_response(
350        StatusCode::INTERNAL_SERVER_ERROR,
351        "api_error",
352        "auth proxy clock error",
353    )
354}
355
356/// Whether the request carries the admin key as a bearer token.
357fn admin_authorized(state: &AuthProxyState, headers: &HeaderMap) -> bool {
358    headers
359        .get(AUTHORIZATION)
360        .and_then(|value| value.to_str().ok())
361        .and_then(|value| value.strip_prefix("Bearer "))
362        .is_some_and(|key| admin_key_matches(&state.config.admin_key, key.trim()))
363}
364
365/// The answer to a pod presenting no valid token. Logs the reason, never the token.
366fn invalid_token(reason: &str, path: &str) -> Response {
367    warn!(reason, path = %path, "request with an invalid token rejected");
368    error_response(
369        StatusCode::UNAUTHORIZED,
370        "authentication_error",
371        "invalid or expired ironflow auth proxy token",
372    )
373}
374
375fn admin_unauthorized() -> Response {
376    warn!("admin request without a valid admin key");
377    error_response(
378        StatusCode::UNAUTHORIZED,
379        "authentication_error",
380        "invalid or missing admin key",
381    )
382}
383
384async fn healthz() -> &'static str {
385    "ok"
386}
387
388async fn issue_token(
389    State(state): State<AuthProxyState>,
390    headers: HeaderMap,
391    body: Bytes,
392) -> Response {
393    if !admin_authorized(&state, &headers) {
394        return admin_unauthorized();
395    }
396    // The body carries a credential: neither the parse error nor the body
397    // is ever echoed or logged.
398    let Ok(request) = from_slice::<TokenRequest>(&body) else {
399        warn!("token request body rejected");
400        return error_response(
401            StatusCode::BAD_REQUEST,
402            "invalid_request_error",
403            "invalid token request body",
404        );
405    };
406    let now = match now_unix() {
407        Ok(now) => now,
408        Err(e) => return clock_error(&e),
409    };
410    let run_id = request.run_id.clone();
411    let step = request.step.clone();
412    match state.registry.issue(request, now) {
413        Ok(issued) => {
414            info!(
415                token = %issued.short_id(),
416                run_id = %run_id,
417                step = %step,
418                "token issued"
419            );
420            (StatusCode::CREATED, Json(issued)).into_response()
421        }
422        Err(AuthProxyError::InvalidRequest(message)) => {
423            warn!(run_id = %run_id, step = %step, reason = %message, "token request refused");
424            error_response(StatusCode::BAD_REQUEST, "invalid_request_error", &message)
425        }
426        Err(e) => {
427            warn!(error = %e, "token issuance failed");
428            error_response(
429                StatusCode::INTERNAL_SERVER_ERROR,
430                "api_error",
431                "token issuance failed",
432            )
433        }
434    }
435}
436
437async fn revoke_token(
438    State(state): State<AuthProxyState>,
439    Path(id): Path<String>,
440    headers: HeaderMap,
441) -> Response {
442    if !admin_authorized(&state, &headers) {
443        return admin_unauthorized();
444    }
445    if state.registry.revoke(&id) {
446        info!(token = %short_id(&id), "token revoked");
447        StatusCode::NO_CONTENT.into_response()
448    } else {
449        error_response(StatusCode::NOT_FOUND, "not_found_error", "unknown token")
450    }
451}
452
453async fn revoke_run(
454    State(state): State<AuthProxyState>,
455    Path(run_id): Path<String>,
456    headers: HeaderMap,
457) -> Response {
458    if !admin_authorized(&state, &headers) {
459        return admin_unauthorized();
460    }
461    let revoked = state.registry.revoke_run(&run_id);
462    info!(run_id = %run_id, revoked, "run tokens revoked");
463    (StatusCode::OK, Json(json!({ "revoked": revoked }))).into_response()
464}
465
466async fn relay(State(state): State<AuthProxyState>, request: Request) -> Response {
467    let (parts, body) = request.into_parts();
468    let method = parts.method;
469    let path = parts.uri.path().to_string();
470
471    // Absolute-form (`GET http://host/..`) and CONNECT ask the proxy to reach
472    // another host.
473    if parts.uri.authority().is_some() || method == Method::CONNECT {
474        warn!(method = %method, "request for another host refused");
475        return error_response(
476            StatusCode::FORBIDDEN,
477            "permission_error",
478            "only api.anthropic.com is reachable through this proxy",
479        );
480    }
481
482    let now = match now_unix() {
483        Ok(now) => now,
484        Err(e) => return clock_error(&e),
485    };
486    let Some(token) = extract_opaque_token(&parts.headers) else {
487        return invalid_token("missing", &path);
488    };
489    let grant = match state.registry.resolve(&token, now) {
490        Ok(grant) => grant,
491        Err(TokenRejection::Unknown) => return invalid_token("unknown", &path),
492        Err(TokenRejection::Expired) => return invalid_token("expired", &path),
493    };
494    let token = short_id(&grant.id);
495
496    if !is_allowed_path(&path) {
497        warn!(token = %token, path = %path, "path outside the API refused");
498        return error_response(
499            StatusCode::FORBIDDEN,
500            "permission_error",
501            "only the Anthropic API under /v1/ is reachable through this proxy",
502        );
503    }
504    if !is_allowed_method(&method) {
505        warn!(token = %token, method = %method, path = %path, "method refused");
506        return error_response(
507            StatusCode::METHOD_NOT_ALLOWED,
508            "invalid_request_error",
509            "only GET and POST are relayed",
510        );
511    }
512
513    let bytes = match to_bytes(body, state.config.max_body_bytes).await {
514        Ok(bytes) => bytes,
515        Err(e) => {
516            warn!(token = %token, error = %e, "request body rejected");
517            return error_response(
518                StatusCode::PAYLOAD_TOO_LARGE,
519                "request_too_large",
520                "request body too large or unreadable",
521            );
522        }
523    };
524
525    let mut url = state.config.upstream.clone();
526    url.set_path(&path);
527    url.set_query(parts.uri.query());
528    let upstream = state
529        .http
530        .request(method.clone(), url)
531        .headers(upstream_headers(&parts.headers, &grant.credential))
532        .body(bytes)
533        .send()
534        .await;
535    let upstream = match upstream {
536        Ok(upstream) => upstream,
537        Err(e) => {
538            warn!(token = %token, error = %e.without_url(), "upstream request failed");
539            return error_response(
540                StatusCode::BAD_GATEWAY,
541                "api_error",
542                "the Anthropic API is unreachable",
543            );
544        }
545    };
546
547    let status = upstream.status();
548    info!(
549        token = %token,
550        run_id = %grant.run_id,
551        step = %grant.step,
552        method = %method,
553        path = %path,
554        status = status.as_u16(),
555        "relayed"
556    );
557    let headers = downstream_headers(upstream.headers());
558    let mut response = Response::new(Body::from_stream(upstream.bytes_stream()));
559    *response.status_mut() = status;
560    *response.headers_mut() = headers;
561    response
562}
ironflow_auth_proxy/lib.rs

ironflow_auth_proxy/
lib.rs
