orchestral_cli/remote/
auth.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::time::{Duration, Instant};
4
5use anyhow::{bail, Context};
6use async_trait::async_trait;
7use axum::http::HeaderName;
8use jsonwebtoken::jwk::JwkSet;
9use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
10use serde_json::Value;
11use tokio::sync::RwLock;
12
13const DEFAULT_JWKS_TTL: Duration = Duration::from_secs(60 * 60);
14
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct GatewayPrincipal {
17    pub subject: Option<String>,
18    pub attributes: BTreeMap<String, String>,
19}
20
21#[derive(Debug, thiserror::Error)]
22pub enum GatewayAuthError {
23    #[error("gateway JWT header is missing")]
24    Missing,
25    #[error("gateway JWT is invalid: {0}")]
26    Invalid(String),
27    #[error("gateway signing keys are unavailable: {0}")]
28    KeysUnavailable(String),
29}
30
31#[async_trait]
32pub trait GatewayAuthenticator: Send + Sync {
33    fn header_name(&self) -> &HeaderName;
34
35    async fn authenticate(&self, token: &str) -> Result<GatewayPrincipal, GatewayAuthError>;
36}
37
38#[derive(Debug, Clone)]
39pub struct JwtGatewayConfig {
40    pub issuer: String,
41    pub audience: String,
42    pub jwks_url: String,
43    pub header_name: HeaderName,
44    pub required_claims: BTreeMap<String, String>,
45}
46
47impl JwtGatewayConfig {
48    pub fn new(
49        issuer: impl Into<String>,
50        audience: impl Into<String>,
51        jwks_url: impl Into<String>,
52        header_name: impl AsRef<str>,
53        required_claims: BTreeMap<String, String>,
54    ) -> anyhow::Result<Self> {
55        let issuer = issuer.into().trim_end_matches('/').to_owned();
56        let audience = audience.into();
57        let jwks_url = jwks_url.into();
58        if issuer.is_empty() || audience.trim().is_empty() || jwks_url.trim().is_empty() {
59            bail!("gateway JWT issuer, audience, and JWKS URL must not be empty");
60        }
61        validate_endpoint("issuer", &issuer)?;
62        validate_endpoint("JWKS URL", &jwks_url)?;
63        let header_name = HeaderName::from_bytes(header_name.as_ref().as_bytes())
64            .context("gateway JWT header name is invalid")?;
65        if required_claims.is_empty() {
66            bail!("at least one --access-jwt-required-claim is required");
67        }
68        if required_claims
69            .iter()
70            .any(|(name, value)| name.trim().is_empty() || value.trim().is_empty())
71        {
72            bail!("gateway JWT required claims must use non-empty NAME=VALUE pairs");
73        }
74        Ok(Self {
75            issuer,
76            audience,
77            jwks_url,
78            header_name,
79            required_claims,
80        })
81    }
82}
83
84fn validate_endpoint(label: &str, value: &str) -> anyhow::Result<()> {
85    let url =
86        reqwest::Url::parse(value).with_context(|| format!("gateway JWT {label} is invalid"))?;
87    let secure = url.scheme() == "https";
88    let local = url.scheme() == "http"
89        && url.host_str().is_some_and(|host| {
90            host == "localhost"
91                || host
92                    .parse::<std::net::IpAddr>()
93                    .is_ok_and(|ip| ip.is_loopback())
94        });
95    if !secure && !local {
96        bail!("gateway JWT {label} must use HTTPS (HTTP is allowed only on loopback)");
97    }
98    Ok(())
99}
100
101#[derive(Debug, Clone)]
102struct CachedKeys {
103    keys: JwkSet,
104    fetched_at: Instant,
105}
106
107#[derive(Debug)]
108pub struct JwtGatewayAuthenticator {
109    config: JwtGatewayConfig,
110    client: reqwest::Client,
111    cache: Arc<RwLock<Option<CachedKeys>>>,
112    cache_ttl: Duration,
113}
114
115impl JwtGatewayAuthenticator {
116    pub fn new(config: JwtGatewayConfig) -> anyhow::Result<Self> {
117        let client = reqwest::Client::builder()
118            .timeout(Duration::from_secs(10))
119            .redirect(reqwest::redirect::Policy::none())
120            .build()
121            .context("build gateway JWT JWKS client")?;
122        Ok(Self {
123            config,
124            client,
125            cache: Arc::new(RwLock::new(None)),
126            cache_ttl: DEFAULT_JWKS_TTL,
127        })
128    }
129
130    async fn keys(&self, force_refresh: bool) -> Result<JwkSet, GatewayAuthError> {
131        if !force_refresh {
132            let cache = self.cache.read().await;
133            if let Some(cached) = cache
134                .as_ref()
135                .filter(|cached| cached.fetched_at.elapsed() < self.cache_ttl)
136            {
137                return Ok(cached.keys.clone());
138            }
139        }
140        let keys = self
141            .client
142            .get(&self.config.jwks_url)
143            .send()
144            .await
145            .map_err(|error| GatewayAuthError::KeysUnavailable(error.to_string()))?
146            .error_for_status()
147            .map_err(|error| GatewayAuthError::KeysUnavailable(error.to_string()))?
148            .json::<JwkSet>()
149            .await
150            .map_err(|error| GatewayAuthError::KeysUnavailable(error.to_string()))?;
151        if keys.keys.is_empty() {
152            return Err(GatewayAuthError::KeysUnavailable(
153                "JWKS endpoint returned no signing keys".to_owned(),
154            ));
155        }
156        *self.cache.write().await = Some(CachedKeys {
157            keys: keys.clone(),
158            fetched_at: Instant::now(),
159        });
160        Ok(keys)
161    }
162
163    async fn decoding_key(&self, kid: &str) -> Result<DecodingKey, GatewayAuthError> {
164        let cached = self.keys(false).await?;
165        if let Some(jwk) = cached.find(kid) {
166            return DecodingKey::from_jwk(jwk)
167                .map_err(|error| GatewayAuthError::Invalid(error.to_string()));
168        }
169        let refreshed = self.keys(true).await?;
170        let jwk = refreshed.find(kid).ok_or_else(|| {
171            GatewayAuthError::Invalid("JWT signing key id was not found in JWKS".to_owned())
172        })?;
173        DecodingKey::from_jwk(jwk).map_err(|error| GatewayAuthError::Invalid(error.to_string()))
174    }
175}
176
177#[async_trait]
178impl GatewayAuthenticator for JwtGatewayAuthenticator {
179    fn header_name(&self) -> &HeaderName {
180        &self.config.header_name
181    }
182
183    async fn authenticate(&self, token: &str) -> Result<GatewayPrincipal, GatewayAuthError> {
184        let header =
185            decode_header(token).map_err(|error| GatewayAuthError::Invalid(error.to_string()))?;
186        if header.alg != Algorithm::RS256 {
187            return Err(GatewayAuthError::Invalid(
188                "only RS256 gateway assertions are accepted".to_owned(),
189            ));
190        }
191        let kid = header.kid.as_deref().ok_or_else(|| {
192            GatewayAuthError::Invalid("JWT header does not contain a key id".to_owned())
193        })?;
194        let key = self.decoding_key(kid).await?;
195        let mut validation = Validation::new(Algorithm::RS256);
196        validation.set_audience(&[&self.config.audience]);
197        validation.set_issuer(&[&self.config.issuer]);
198        validation.set_required_spec_claims(&["exp", "aud", "iss"]);
199        validation.validate_nbf = true;
200        let claims = decode::<Value>(token, &key, &validation)
201            .map_err(|error| GatewayAuthError::Invalid(error.to_string()))?
202            .claims;
203
204        let mut attributes = BTreeMap::new();
205        for (name, expected) in &self.config.required_claims {
206            let value = claim_at_path(&claims, name).ok_or_else(|| {
207                GatewayAuthError::Invalid(format!("required claim '{name}' is missing"))
208            })?;
209            if !claim_contains(value, expected) {
210                return Err(GatewayAuthError::Invalid(format!(
211                    "required claim '{name}' does not match"
212                )));
213            }
214            attributes.insert(name.clone(), expected.clone());
215        }
216        Ok(GatewayPrincipal {
217            subject: claims.get("sub").and_then(Value::as_str).map(str::to_owned),
218            attributes,
219        })
220    }
221}
222
223fn claim_at_path<'a>(claims: &'a Value, path: &str) -> Option<&'a Value> {
224    path.split('.')
225        .try_fold(claims, |value, segment| value.get(segment))
226}
227
228fn claim_contains(value: &Value, expected: &str) -> bool {
229    match value {
230        Value::String(actual) => actual == expected,
231        Value::Array(values) => values.iter().any(|value| claim_contains(value, expected)),
232        Value::Bool(actual) => expected.parse::<bool>() == Ok(*actual),
233        Value::Number(actual) => actual.to_string() == expected,
234        _ => false,
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use serde_json::json;
242
243    #[test]
244    fn required_claims_support_nested_values_and_group_membership() {
245        let claims = json!({"identity": {"email": "person@example.com"}, "groups": ["dev", "ops"]});
246        assert!(claim_contains(
247            claim_at_path(&claims, "identity.email").unwrap(),
248            "person@example.com"
249        ));
250        assert!(claim_contains(
251            claim_at_path(&claims, "groups").unwrap(),
252            "ops"
253        ));
254        assert!(!claim_contains(
255            claim_at_path(&claims, "groups").unwrap(),
256            "admin"
257        ));
258    }
259
260    #[test]
261    fn gateway_endpoints_require_https_except_for_loopback_tests() {
262        let claims = BTreeMap::from([("email".to_owned(), "person@example.com".to_owned())]);
263        assert!(JwtGatewayConfig::new(
264            "https://access.example.com",
265            "audience",
266            "https://access.example.com/keys",
267            "x-access-jwt",
268            claims.clone(),
269        )
270        .is_ok());
271        assert!(JwtGatewayConfig::new(
272            "http://access.example.com",
273            "audience",
274            "http://access.example.com/keys",
275            "x-access-jwt",
276            claims,
277        )
278        .is_err());
279    }
280}
orchestral_cli/remote/auth.rs

orchestral_cli/remote/
auth.rs
