gcloud_sdk/
axum_layer.rs

1//! Protects an axum service with Google ID tokens: [`VerifyIdTokenLayer`] verifies the
2//! bearer token of each request with an [`IdTokenVerifier`], and handlers read the
3//! verified claims with the [`VerifiedIdToken`] extractor.
4//!
5//! The layer is a plain tower layer, so it also wraps a tonic server.
6//!
7//! ```ignore
8//! let verifier = Arc::new(IdTokenVerifier::new(IdTokenAudience::new(
9//!     "https://orders-abc123-ew.a.run.app",
10//! ))?);
11//! let app = Router::new()
12//!     .route("/orders", post(create_order))
13//!     .layer(VerifyIdTokenLayer::new(verifier).authorize(|token| {
14//!         token.verified_email() == Some(&PrincipalEmail::new("caller@my-project.iam.gserviceaccount.com"))
15//!     }));
16//!
17//! async fn create_order(token: VerifiedIdToken) -> String {
18//!     format!("hello, {}", token.subject())
19//! }
20//! ```
21
22use std::pin::Pin;
23use std::sync::Arc;
24use std::task::{Context, Poll};
25
26use axum_core::extract::FromRequestParts;
27use futures::Future;
28use hyper::header::{HeaderMap, HeaderValue, AUTHORIZATION, WWW_AUTHENTICATE};
29use hyper::http::request::Parts;
30use hyper::{Request, Response, StatusCode};
31use tower::Service;
32use tower_layer::Layer;
33use tracing::*;
34
35use crate::id_token_verify::{IdTokenVerifier, IdTokenVerifyError, VerifiedIdToken};
36
37type Authorize = dyn Fn(&VerifiedIdToken) -> bool + Send + Sync;
38
39/// A tower layer that lets a request through only with a valid Google ID token in its
40/// `authorization: Bearer` header, and puts the token's claims into the request's
41/// extensions for the [`VerifiedIdToken`] extractor.
42///
43/// Requests it refuses never reach the inner service. They get an empty response:
44/// - 401 with `www-authenticate: Bearer` when there is no bearer token;
45/// - 401 with `www-authenticate: Bearer error="invalid_token"` when the token does not
46///   verify;
47/// - 503 when the signing keys cannot be fetched;
48/// - 403 when the token verifies but [`authorize`](Self::authorize) refuses it.
49#[derive(Clone)]
50pub struct VerifyIdTokenLayer {
51    verifier: Arc<IdTokenVerifier>,
52    authorize: Option<Arc<Authorize>>,
53}
54
55impl VerifyIdTokenLayer {
56    /// A layer that admits every token `verifier` accepts.
57    pub fn new(verifier: Arc<IdTokenVerifier>) -> Self {
58        Self {
59            verifier,
60            authorize: None,
61        }
62    }
63
64    /// Admits only the verified tokens for which `authorize` returns true, such as those
65    /// of an allowlist of callers; the others get 403. Replaces any earlier `authorize`.
66    pub fn authorize(
67        mut self,
68        authorize: impl Fn(&VerifiedIdToken) -> bool + Send + Sync + 'static,
69    ) -> Self {
70        self.authorize = Some(Arc::new(authorize));
71        self
72    }
73
74    async fn admit(&self, headers: &HeaderMap) -> Result<VerifiedIdToken, Refusal> {
75        let token = bearer_token(headers).ok_or(Refusal::MissingToken)?;
76        let verified = self.verifier.verify(token).await.map_err(|err| match err {
77            IdTokenVerifyError::InvalidToken(reason) => {
78                debug!(%reason, "Refusing a request with an invalid ID token");
79                Refusal::InvalidToken
80            }
81            IdTokenVerifyError::KeysUnavailable(error) => {
82                warn!(%error, "Refusing a request: ID token signing keys are unavailable");
83                Refusal::KeysUnavailable
84            }
85        })?;
86        match &self.authorize {
87            Some(authorize) if !authorize(&verified) => Err(Refusal::Forbidden),
88            _ => Ok(verified),
89        }
90    }
91}
92
93impl<S> Layer<S> for VerifyIdTokenLayer {
94    type Service = VerifyIdToken<S>;
95
96    fn layer(&self, inner: S) -> Self::Service {
97        VerifyIdToken {
98            inner,
99            layer: self.clone(),
100        }
101    }
102}
103
104/// The service made by [`VerifyIdTokenLayer`].
105#[derive(Clone)]
106pub struct VerifyIdToken<S> {
107    inner: S,
108    layer: VerifyIdTokenLayer,
109}
110
111impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for VerifyIdToken<S>
112where
113    S: Service<Request<ReqBody>, Response = Response<ResBody>> + Clone + Send + 'static,
114    S::Future: Send + 'static,
115    ReqBody: Send + 'static,
116    ResBody: Default + Send + 'static,
117{
118    type Response = Response<ResBody>;
119    type Error = S::Error;
120    type Future =
121        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
122
123    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
124        self.inner.poll_ready(cx)
125    }
126
127    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
128        let layer = self.layer.clone();
129        // tower's idiom for a `Clone` inner service: the instance polled ready goes into
130        // the future, and the service keeps a fresh clone.
131        let clone = self.inner.clone();
132        let mut inner = std::mem::replace(&mut self.inner, clone);
133
134        Box::pin(async move {
135            match layer.admit(req.headers()).await {
136                Ok(verified) => {
137                    req.extensions_mut().insert(verified);
138                    inner.call(req).await
139                }
140                Err(refusal) => Ok(refusal.into()),
141            }
142        })
143    }
144}
145
146/// Why a request was refused.
147enum Refusal {
148    MissingToken,
149    InvalidToken,
150    KeysUnavailable,
151    Forbidden,
152}
153
154impl<B: Default> From<Refusal> for Response<B> {
155    fn from(refusal: Refusal) -> Self {
156        let (status, challenge) = match refusal {
157            Refusal::MissingToken => (StatusCode::UNAUTHORIZED, Some("Bearer")),
158            Refusal::InvalidToken => (
159                StatusCode::UNAUTHORIZED,
160                Some(r#"Bearer error="invalid_token""#),
161            ),
162            Refusal::KeysUnavailable => (StatusCode::SERVICE_UNAVAILABLE, None),
163            Refusal::Forbidden => (StatusCode::FORBIDDEN, None),
164        };
165        let mut response = Response::new(B::default());
166        *response.status_mut() = status;
167        if let Some(challenge) = challenge {
168            response
169                .headers_mut()
170                .insert(WWW_AUTHENTICATE, HeaderValue::from_static(challenge));
171        }
172        response
173    }
174}
175
176/// The token of an `authorization: Bearer <token>` header, with the scheme matched
177/// case-insensitively.
178fn bearer_token(headers: &HeaderMap) -> Option<&str> {
179    let (scheme, token) = headers.get(AUTHORIZATION)?.to_str().ok()?.split_once(' ')?;
180    let token = token.trim();
181    (scheme.eq_ignore_ascii_case("bearer") && !token.is_empty()).then_some(token)
182}
183
184/// The claims of the token [`VerifyIdTokenLayer`] verified for this request.
185///
186/// Rejects with 500 when the route is not behind a [`VerifyIdTokenLayer`], since that
187/// is a mistake in the server, not in the request.
188impl<S: Send + Sync> FromRequestParts<S> for VerifiedIdToken {
189    type Rejection = StatusCode;
190
191    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
192        parts
193            .extensions
194            .get::<VerifiedIdToken>()
195            .cloned()
196            .ok_or(StatusCode::INTERNAL_SERVER_ERROR)
197    }
198}
199
200// The tests sign tokens, which needs a crypto provider.
201#[cfg(all(test, any(feature = "jwt-aws-lc-rs", feature = "jwt-rust-crypto")))]
202mod tests {
203    use super::*;
204    use crate::error::ErrorKind;
205    use crate::id_token_verify::{IdTokenKeys, IdTokenKeysSource, PrincipalEmail};
206    use crate::test_support::{signed_jwt, test_jwk_set};
207    use crate::IdTokenAudience;
208    use async_trait::async_trait;
209    use jiff::{SignedDuration, Timestamp};
210    use std::convert::Infallible;
211    use std::time::Duration;
212
213    const AUDIENCE: &str = "https://orders-abc123-ew.a.run.app";
214    const CALLER: &str = "caller@my-project.iam.gserviceaccount.com";
215    const KEY_ID: &str = "current-key";
216
217    enum TestKeys {
218        Published,
219        Unavailable,
220    }
221
222    #[async_trait]
223    impl IdTokenKeysSource for TestKeys {
224        async fn fetch_keys(&self) -> crate::error::Result<IdTokenKeys> {
225            match self {
226                Self::Published => Ok(IdTokenKeys::new(
227                    test_jwk_set(&[KEY_ID]),
228                    Duration::from_secs(3600),
229                )),
230                Self::Unavailable => {
231                    Err(ErrorKind::HttpStatus(reqwest::StatusCode::SERVICE_UNAVAILABLE).into())
232                }
233            }
234        }
235    }
236
237    fn verifier(keys: TestKeys) -> Arc<IdTokenVerifier> {
238        Arc::new(IdTokenVerifier::with_keys_source(
239            IdTokenAudience::new(AUDIENCE),
240            keys,
241        ))
242    }
243
244    fn token_for(audience: &str) -> String {
245        let now = Timestamp::now();
246        let claims = serde_json::json!({
247            "iss": "https://accounts.google.com",
248            "aud": audience,
249            "sub": "112233445566778899",
250            "email": CALLER,
251            "email_verified": true,
252            "iat": now.as_second(),
253            "exp": (now + SignedDuration::from_hours(1)).as_second(),
254        });
255        signed_jwt(KEY_ID, &claims)
256    }
257
258    fn request(authorization: Option<String>) -> Request<String> {
259        let mut builder = Request::builder().uri("http://orders.internal/orders");
260        if let Some(authorization) = authorization {
261            builder = builder.header(AUTHORIZATION, authorization);
262        }
263        builder.body(String::new()).unwrap()
264    }
265
266    /// Answers 200 and copies the claims the layer passed on into the response's
267    /// extensions, so that tests see what the inner service saw.
268    #[derive(Clone)]
269    struct EchoClaims;
270
271    impl Service<Request<String>> for EchoClaims {
272        type Response = Response<String>;
273        type Error = Infallible;
274        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
275
276        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
277            Poll::Ready(Ok(()))
278        }
279
280        fn call(&mut self, req: Request<String>) -> Self::Future {
281            let claims = req.extensions().get::<VerifiedIdToken>().cloned();
282            Box::pin(async move {
283                let mut response = Response::new(String::new());
284                if let Some(claims) = claims {
285                    response.extensions_mut().insert(claims);
286                }
287                Ok(response)
288            })
289        }
290    }
291
292    async fn send(layer: VerifyIdTokenLayer, req: Request<String>) -> Response<String> {
293        layer.layer(EchoClaims).call(req).await.unwrap()
294    }
295
296    #[tokio::test]
297    async fn missing_token_is_unauthorized() {
298        let response = send(
299            VerifyIdTokenLayer::new(verifier(TestKeys::Published)),
300            request(None),
301        )
302        .await;
303
304        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
305        assert_eq!(
306            response.headers().get(WWW_AUTHENTICATE),
307            Some(&HeaderValue::from_static("Bearer"))
308        );
309        assert!(response.extensions().get::<VerifiedIdToken>().is_none());
310    }
311
312    #[tokio::test]
313    async fn token_for_another_audience_is_unauthorized() {
314        let token = token_for("https://billing-abc123-ew.a.run.app");
315
316        let response = send(
317            VerifyIdTokenLayer::new(verifier(TestKeys::Published)),
318            request(Some(format!("Bearer {token}"))),
319        )
320        .await;
321
322        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
323        assert_eq!(
324            response.headers().get(WWW_AUTHENTICATE),
325            Some(&HeaderValue::from_static(r#"Bearer error="invalid_token""#))
326        );
327        assert!(response.extensions().get::<VerifiedIdToken>().is_none());
328    }
329
330    #[tokio::test]
331    async fn unavailable_keys_are_service_unavailable() {
332        let response = send(
333            VerifyIdTokenLayer::new(verifier(TestKeys::Unavailable)),
334            request(Some(format!("Bearer {}", token_for(AUDIENCE)))),
335        )
336        .await;
337
338        assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
339        assert!(response.extensions().get::<VerifiedIdToken>().is_none());
340    }
341
342    #[tokio::test]
343    async fn refused_principal_is_forbidden() {
344        let layer = VerifyIdTokenLayer::new(verifier(TestKeys::Published)).authorize(|token| {
345            token.verified_email()
346                == Some(&PrincipalEmail::new(
347                    "billing@my-project.iam.gserviceaccount.com",
348                ))
349        });
350
351        let response = send(
352            layer,
353            request(Some(format!("Bearer {}", token_for(AUDIENCE)))),
354        )
355        .await;
356
357        assert_eq!(response.status(), StatusCode::FORBIDDEN);
358        assert!(response.extensions().get::<VerifiedIdToken>().is_none());
359    }
360
361    #[tokio::test]
362    async fn verified_claims_reach_the_inner_service() {
363        let layer = VerifyIdTokenLayer::new(verifier(TestKeys::Published))
364            .authorize(|token| token.verified_email() == Some(&PrincipalEmail::new(CALLER)));
365
366        let response = send(
367            layer,
368            request(Some(format!("bearer {}", token_for(AUDIENCE)))),
369        )
370        .await;
371
372        assert_eq!(response.status(), StatusCode::OK);
373        let claims = response.extensions().get::<VerifiedIdToken>().unwrap();
374        assert_eq!(claims.verified_email(), Some(&PrincipalEmail::new(CALLER)));
375    }
376
377    #[tokio::test]
378    async fn extractor_reads_the_verified_claims() {
379        let claims = verifier(TestKeys::Published)
380            .verify(&token_for(AUDIENCE))
381            .await
382            .unwrap();
383        let (mut parts, _) = request(None).into_parts();
384        parts.extensions.insert(claims);
385
386        let extracted = VerifiedIdToken::from_request_parts(&mut parts, &())
387            .await
388            .unwrap();
389
390        assert_eq!(
391            extracted.verified_email(),
392            Some(&PrincipalEmail::new(CALLER))
393        );
394    }
395
396    #[tokio::test]
397    async fn extractor_outside_the_layer_is_an_internal_error() {
398        let (mut parts, _) = request(None).into_parts();
399
400        let rejection = VerifiedIdToken::from_request_parts(&mut parts, &())
401            .await
402            .unwrap_err();
403
404        assert_eq!(rejection, StatusCode::INTERNAL_SERVER_ERROR);
405    }
406}
gcloud_sdk/axum_layer.rs

gcloud_sdk/
axum_layer.rs
