gcloud_sdk/
middleware.rs

1use crate::token_source::auth_token_generator::GoogleAuthTokenGenerator;
2use futures::{Future, TryFutureExt};
3use hyper::header::{HeaderMap, HeaderName, HeaderValue, USER_AGENT};
4use jiff::Timestamp;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context, Poll};
8use tonic::client::GrpcService;
9use tower::Service;
10use tower_layer::Layer;
11use tracing::*;
12
13const X_GOOG_API_CLIENT: HeaderName = HeaderName::from_static("x-goog-api-client");
14const GOOGLE_CLOUD_RESOURCE_PREFIX: HeaderName =
15    HeaderName::from_static("google-cloud-resource-prefix");
16
17fn default_headers(cloud_resource_prefix: Option<String>) -> crate::error::Result<HeaderMap> {
18    let mut headers = HeaderMap::new();
19    let default_agent =
20        HeaderValue::from_static(concat!("gcloud-sdk-rs/", env!("CARGO_PKG_VERSION")));
21    headers.insert(USER_AGENT, default_agent.clone());
22    headers.insert(X_GOOG_API_CLIENT, default_agent);
23    if let Some(prefix) = cloud_resource_prefix {
24        headers.insert(
25            GOOGLE_CLOUD_RESOURCE_PREFIX,
26            HeaderValue::from_str(&prefix)?,
27        );
28    }
29    Ok(headers)
30}
31
32/// Appends `extra` to whatever `name` currently holds in `headers` (space separated),
33/// or sets it outright if absent.
34fn append_header_value(
35    headers: &HeaderMap,
36    name: &HeaderName,
37    extra: &str,
38) -> crate::error::Result<HeaderValue> {
39    let combined = match headers.get(name).and_then(|v| v.to_str().ok()) {
40        Some(current) => format!("{current} {extra}"),
41        None => extra.to_string(),
42    };
43    Ok(HeaderValue::from_str(&combined)?)
44}
45
46#[derive(Clone)]
47pub struct GoogleAuthMiddlewareService<T> {
48    inner: T,
49    token_generator: Arc<GoogleAuthTokenGenerator>,
50    /// Every header added to each request except `authorization`, already validated.
51    headers: Arc<HeaderMap>,
52}
53
54impl<T> GoogleAuthMiddlewareService<T> {
55    pub fn new(
56        service: T,
57        token_generator: Arc<GoogleAuthTokenGenerator>,
58        cloud_resource_prefix: Option<String>,
59    ) -> crate::error::Result<GoogleAuthMiddlewareService<T>> {
60        Ok(GoogleAuthMiddlewareService {
61            inner: service,
62            token_generator,
63            headers: Arc::new(default_headers(cloud_resource_prefix)?),
64        })
65    }
66
67    pub fn set_user_agent(&mut self, user_agent: String) -> crate::error::Result<()> {
68        let value = HeaderValue::from_str(&user_agent)?;
69        Arc::make_mut(&mut self.headers).insert(USER_AGENT, value);
70        Ok(())
71    }
72
73    pub fn set_x_goog_api_client(&mut self, x_goog_api_client: String) -> crate::error::Result<()> {
74        let value = HeaderValue::from_str(&x_goog_api_client)?;
75        Arc::make_mut(&mut self.headers).insert(X_GOOG_API_CLIENT, value);
76        Ok(())
77    }
78
79    pub fn set_cloud_resource_prefix(
80        &mut self,
81        cloud_resource_prefix: String,
82    ) -> crate::error::Result<()> {
83        let value = HeaderValue::from_str(&cloud_resource_prefix)?;
84        Arc::make_mut(&mut self.headers).insert(GOOGLE_CLOUD_RESOURCE_PREFIX, value);
85        Ok(())
86    }
87
88    pub fn append_user_agent(&mut self, user_agent: String) -> crate::error::Result<()> {
89        let value = append_header_value(&self.headers, &USER_AGENT, &user_agent)?;
90        Arc::make_mut(&mut self.headers).insert(USER_AGENT, value);
91        Ok(())
92    }
93
94    pub fn append_x_goog_api_client(
95        &mut self,
96        x_goog_api_client: String,
97    ) -> crate::error::Result<()> {
98        let value = append_header_value(&self.headers, &X_GOOG_API_CLIENT, &x_goog_api_client)?;
99        Arc::make_mut(&mut self.headers).insert(X_GOOG_API_CLIENT, value);
100        Ok(())
101    }
102
103    pub fn set_additional_headers(&mut self, additional_headers: HeaderMap) {
104        Arc::make_mut(&mut self.headers).extend(additional_headers);
105    }
106
107    /// Wraps `inner` in a middleware service that shares this one's token generator, so
108    /// both fetch and refresh one token, and carries the same headers except
109    /// `google-cloud-resource-prefix`, which is set to `cloud_resource_prefix` or left
110    /// out when it is `None`.
111    pub(crate) fn with_inner<U>(
112        &self,
113        inner: U,
114        cloud_resource_prefix: Option<String>,
115    ) -> crate::error::Result<GoogleAuthMiddlewareService<U>> {
116        let mut headers = HeaderMap::clone(&self.headers);
117        headers.remove(GOOGLE_CLOUD_RESOURCE_PREFIX);
118        if let Some(prefix) = cloud_resource_prefix {
119            headers.insert(
120                GOOGLE_CLOUD_RESOURCE_PREFIX,
121                HeaderValue::from_str(&prefix)?,
122            );
123        }
124        Ok(GoogleAuthMiddlewareService {
125            inner,
126            token_generator: Arc::clone(&self.token_generator),
127            headers: Arc::new(headers),
128        })
129    }
130}
131
132impl<T, RequestBody> Service<hyper::Request<RequestBody>> for GoogleAuthMiddlewareService<T>
133where
134    T: GrpcService<RequestBody> + Send + Clone + 'static,
135    T::Future: 'static + Send,
136    RequestBody: 'static + Send,
137    T::ResponseBody: 'static + Send,
138    T::Error: 'static + Send,
139{
140    type Response = hyper::Response<T::ResponseBody>;
141    type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
142    type Future =
143        Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
144
145    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
146        self.inner.poll_ready(cx).map_err(Into::into)
147    }
148
149    fn call(&mut self, mut req: hyper::Request<RequestBody>) -> Self::Future {
150        let generator = Arc::clone(&self.token_generator);
151        let headers = Arc::clone(&self.headers);
152
153        // tower's documented idiom for a `Clone` inner service: the instance we already
154        // polled ready goes into the future, and the service keeps a fresh clone.
155        let clone = self.inner.clone();
156        let mut inner = std::mem::replace(&mut self.inner, clone);
157
158        Box::pin(async move {
159            let begin_time = Timestamp::now();
160            let authorization = generator.authorization_header().await.map_err(Box::new)?;
161            let token_generated_time = Timestamp::now();
162
163            let req_headers = req.headers_mut();
164            req_headers.insert(hyper::header::AUTHORIZATION, authorization);
165            // Each name in `headers` fully replaces whatever the request already carries
166            // under it (multi-valued or not); `iter()` yields one pair per value, so the
167            // names are cleared first and then appended rather than repeatedly `insert`ed,
168            // which would silently drop every value but the last for a repeated name.
169            for name in headers.keys() {
170                req_headers.remove(name);
171            }
172            req_headers.extend(
173                headers
174                    .iter()
175                    .map(|(name, value)| (name.clone(), value.clone())),
176            );
177
178            let req_uri = req.uri().clone();
179            inner
180                .call(req)
181                .map_ok(|x| {
182                    let finished_time = Timestamp::now();
183                    debug!(
184                        %req_uri,
185                        "OK: took {}ms (incl. token gen: {}ms)",
186                        finished_time.duration_since(begin_time).as_millis(),
187                        token_generated_time.duration_since(begin_time).as_millis()
188                    );
189                    x
190                })
191                .await
192                .map_err(|e| {
193                    let finished_time = Timestamp::now();
194                    error!(
195                        %req_uri,
196                        "Err: took {}ms (incl. token gen: {}ms)",
197                        finished_time.duration_since(begin_time).as_millis(),
198                        token_generated_time.duration_since(begin_time).as_millis()
199                    );
200                    e.into()
201                })
202        })
203    }
204}
205
206pub struct GoogleAuthMiddlewareLayer {
207    token_generator: Arc<GoogleAuthTokenGenerator>,
208    headers: Arc<HeaderMap>,
209}
210
211impl GoogleAuthMiddlewareLayer {
212    pub fn new(
213        token_generator: GoogleAuthTokenGenerator,
214        cloud_resource_prefix: Option<String>,
215    ) -> crate::error::Result<Self> {
216        Ok(GoogleAuthMiddlewareLayer {
217            token_generator: Arc::new(token_generator),
218            headers: Arc::new(default_headers(cloud_resource_prefix)?),
219        })
220    }
221
222    pub fn amend_user_agent(mut self, user_agent: String) -> crate::error::Result<Self> {
223        let value = append_header_value(&self.headers, &USER_AGENT, &user_agent)?;
224        Arc::make_mut(&mut self.headers).insert(USER_AGENT, value);
225        Ok(self)
226    }
227
228    pub fn amend_x_goog_api_client(
229        mut self,
230        x_goog_api_client: String,
231    ) -> crate::error::Result<Self> {
232        let value = append_header_value(&self.headers, &X_GOOG_API_CLIENT, &x_goog_api_client)?;
233        Arc::make_mut(&mut self.headers).insert(X_GOOG_API_CLIENT, value);
234        Ok(self)
235    }
236
237    pub fn set_additional_headers(&mut self, additional_headers: HeaderMap) {
238        Arc::make_mut(&mut self.headers).extend(additional_headers);
239    }
240}
241
242impl<S> Layer<S> for GoogleAuthMiddlewareLayer {
243    type Service = GoogleAuthMiddlewareService<S>;
244
245    fn layer(&self, service: S) -> GoogleAuthMiddlewareService<S> {
246        GoogleAuthMiddlewareService {
247            inner: service,
248            token_generator: Arc::clone(&self.token_generator),
249            headers: Arc::clone(&self.headers),
250        }
251    }
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::token_source::{Source, Token, TokenSourceType};
258    use async_trait::async_trait;
259    use hyper::{Request, Response};
260    use jiff::{SignedDuration, Timestamp};
261    use secret_vault_value::SecretValue;
262    use std::convert::Infallible;
263
264    struct DummySource;
265
266    #[async_trait]
267    impl Source for DummySource {
268        async fn token(&self) -> crate::error::Result<Token> {
269            Ok(Token {
270                token_type: "Bearer".to_string(),
271                token: SecretValue::from("dummy-token"),
272                expiry: Timestamp::now() + SignedDuration::from_hours(1),
273            })
274        }
275    }
276
277    struct CountingSource {
278        calls: Arc<std::sync::atomic::AtomicUsize>,
279    }
280
281    #[async_trait]
282    impl Source for CountingSource {
283        async fn token(&self) -> crate::error::Result<Token> {
284            self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
285            Ok(Token {
286                token_type: "Bearer".to_string(),
287                token: SecretValue::from("counted-token"),
288                expiry: Timestamp::now() + SignedDuration::from_hours(1),
289            })
290        }
291    }
292
293    #[derive(Clone)]
294    struct DummyService {
295        tx: Arc<tokio::sync::mpsc::Sender<Request<String>>>,
296    }
297
298    impl Service<Request<String>> for DummyService {
299        type Response = Response<String>;
300        type Error = Infallible;
301        type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
302
303        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
304            Poll::Ready(Ok(()))
305        }
306
307        fn call(&mut self, req: Request<String>) -> Self::Future {
308            let tx = self.tx.clone();
309            Box::pin(async move {
310                tx.send(req).await.unwrap();
311                Ok(Response::builder()
312                    .status(200)
313                    .body("".to_string())
314                    .unwrap())
315            })
316        }
317    }
318
319    #[tokio::test]
320    async fn test_headers_presence() {
321        let token_generator = GoogleAuthTokenGenerator::new(
322            TokenSourceType::ExternalSource(Box::new(DummySource)),
323            vec![],
324        )
325        .await
326        .unwrap();
327
328        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
329        let dummy_service = DummyService { tx: Arc::new(tx) };
330        let mut service =
331            GoogleAuthMiddlewareService::new(dummy_service, Arc::new(token_generator), None)
332                .unwrap();
333
334        let req = Request::builder()
335            .uri("http://example.com")
336            .body("".to_string())
337            .unwrap();
338
339        tower::Service::call(&mut service, req).await.unwrap();
340
341        let captured_req = rx.recv().await.unwrap();
342        let expected_default = format!("gcloud-sdk-rs/{}", env!("CARGO_PKG_VERSION"));
343        assert_eq!(
344            captured_req
345                .headers()
346                .get(hyper::header::USER_AGENT)
347                .unwrap(),
348            expected_default.as_str()
349        );
350        assert_eq!(
351            captured_req.headers().get("x-goog-api-client").unwrap(),
352            expected_default.as_str()
353        );
354        assert_eq!(
355            captured_req.headers().get("authorization").unwrap(),
356            "Bearer dummy-token"
357        );
358    }
359
360    #[tokio::test]
361    async fn authorization_header_is_marked_sensitive() {
362        let token_generator = GoogleAuthTokenGenerator::new(
363            TokenSourceType::ExternalSource(Box::new(DummySource)),
364            vec![],
365        )
366        .await
367        .unwrap();
368
369        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
370        let dummy_service = DummyService { tx: Arc::new(tx) };
371        let mut service =
372            GoogleAuthMiddlewareService::new(dummy_service, Arc::new(token_generator), None)
373                .unwrap();
374
375        let req = Request::builder()
376            .uri("http://example.com")
377            .body("".to_string())
378            .unwrap();
379
380        tower::Service::call(&mut service, req).await.unwrap();
381
382        let captured_req = rx.recv().await.unwrap();
383        assert!(captured_req
384            .headers()
385            .get("authorization")
386            .unwrap()
387            .is_sensitive());
388    }
389
390    #[tokio::test]
391    async fn test_headers_amend() {
392        let token_generator = GoogleAuthTokenGenerator::new(
393            TokenSourceType::ExternalSource(Box::new(DummySource)),
394            vec![],
395        )
396        .await
397        .unwrap();
398
399        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
400        let dummy_service = DummyService { tx: Arc::new(tx) };
401
402        let layer = GoogleAuthMiddlewareLayer::new(token_generator, None)
403            .unwrap()
404            .amend_user_agent("extra-ua".to_string())
405            .unwrap()
406            .amend_x_goog_api_client("extra-client".to_string())
407            .unwrap();
408
409        let mut service = layer.layer(dummy_service);
410
411        let req = Request::builder()
412            .uri("http://example.com")
413            .body("".to_string())
414            .unwrap();
415
416        tower::Service::call(&mut service, req).await.unwrap();
417
418        let captured_req = rx.recv().await.unwrap();
419        let expected_ua = format!("gcloud-sdk-rs/{} extra-ua", env!("CARGO_PKG_VERSION"));
420        let expected_client = format!("gcloud-sdk-rs/{} extra-client", env!("CARGO_PKG_VERSION"));
421
422        assert_eq!(
423            captured_req
424                .headers()
425                .get(hyper::header::USER_AGENT)
426                .unwrap(),
427            expected_ua.as_str()
428        );
429        assert_eq!(
430            captured_req.headers().get("x-goog-api-client").unwrap(),
431            expected_client.as_str()
432        );
433    }
434
435    #[tokio::test]
436    async fn invalid_user_agent_is_rejected_at_setter() {
437        let token_generator = GoogleAuthTokenGenerator::new(
438            TokenSourceType::ExternalSource(Box::new(DummySource)),
439            vec![],
440        )
441        .await
442        .unwrap();
443
444        let layer_result = GoogleAuthMiddlewareLayer::new(token_generator, None)
445            .unwrap()
446            .amend_user_agent("bad\nvalue".to_string());
447
448        match layer_result {
449            Err(e) => assert!(matches!(
450                e.into_kind(),
451                crate::error::ErrorKind::HeaderValue(_)
452            )),
453            Ok(_) => panic!("expected an invalid header value to be rejected"),
454        }
455    }
456
457    #[tokio::test]
458    async fn amended_clone_does_not_change_sibling() {
459        let token_generator = GoogleAuthTokenGenerator::new(
460            TokenSourceType::ExternalSource(Box::new(DummySource)),
461            vec![],
462        )
463        .await
464        .unwrap();
465
466        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
467        let dummy_service = DummyService { tx: Arc::new(tx) };
468        let base_service =
469            GoogleAuthMiddlewareService::new(dummy_service, Arc::new(token_generator), None)
470                .unwrap();
471
472        let mut amended = base_service.clone();
473        amended.append_user_agent("extra".to_string()).unwrap();
474
475        let mut sibling = base_service.clone();
476
477        let req = Request::builder()
478            .uri("http://example.com")
479            .body("".to_string())
480            .unwrap();
481
482        tower::Service::call(&mut sibling, req).await.unwrap();
483
484        let captured_req = rx.recv().await.unwrap();
485        let expected_default = format!("gcloud-sdk-rs/{}", env!("CARGO_PKG_VERSION"));
486        assert_eq!(
487            captured_req
488                .headers()
489                .get(hyper::header::USER_AGENT)
490                .unwrap(),
491            expected_default.as_str()
492        );
493    }
494
495    #[tokio::test]
496    async fn test_additional_headers() {
497        let token_generator = GoogleAuthTokenGenerator::new(
498            TokenSourceType::ExternalSource(Box::new(DummySource)),
499            vec![],
500        )
501        .await
502        .unwrap();
503
504        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
505        let dummy_service = DummyService { tx: Arc::new(tx) };
506        let mut service =
507            GoogleAuthMiddlewareService::new(dummy_service, Arc::new(token_generator), None)
508                .unwrap();
509        let mut test_headers = hyper::HeaderMap::new();
510        test_headers.insert("x-test-header", "test-value".parse().unwrap());
511        service.set_additional_headers(test_headers);
512
513        let req = Request::builder()
514            .uri("http://example.com")
515            .body("".to_string())
516            .unwrap();
517
518        tower::Service::call(&mut service, req).await.unwrap();
519
520        let captured_req = rx.recv().await.unwrap();
521        assert_eq!(
522            captured_req.headers().get("x-test-header").unwrap(),
523            "test-value"
524        );
525    }
526
527    #[tokio::test]
528    async fn additional_headers_keep_every_value_of_a_repeated_name() {
529        let token_generator = GoogleAuthTokenGenerator::new(
530            TokenSourceType::ExternalSource(Box::new(DummySource)),
531            vec![],
532        )
533        .await
534        .unwrap();
535
536        let (tx, mut rx) = tokio::sync::mpsc::channel(1);
537        let dummy_service = DummyService { tx: Arc::new(tx) };
538        let mut service =
539            GoogleAuthMiddlewareService::new(dummy_service, Arc::new(token_generator), None)
540                .unwrap();
541
542        let mut test_headers = hyper::HeaderMap::new();
543        test_headers.append("x-multi", "first".parse().unwrap());
544        test_headers.append("x-multi", "second".parse().unwrap());
545        service.set_additional_headers(test_headers);
546
547        let req = Request::builder()
548            .uri("http://example.com")
549            .body("".to_string())
550            .unwrap();
551
552        tower::Service::call(&mut service, req).await.unwrap();
553
554        let captured_req = rx.recv().await.unwrap();
555        let values: Vec<&str> = captured_req
556            .headers()
557            .get_all("x-multi")
558            .iter()
559            .map(|v| v.to_str().unwrap())
560            .collect();
561        assert_eq!(values, vec!["first", "second"]);
562    }
563
564    #[tokio::test]
565    async fn with_inner_shares_the_token_and_replaces_the_resource_prefix() {
566        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
567        let token_generator = GoogleAuthTokenGenerator::new(
568            TokenSourceType::ExternalSource(Box::new(CountingSource {
569                calls: calls.clone(),
570            })),
571            vec![],
572        )
573        .await
574        .unwrap();
575
576        let (tx, mut rx) = tokio::sync::mpsc::channel(2);
577        let mut extra = hyper::HeaderMap::new();
578        extra.insert("x-test-header", "test-value".parse().unwrap());
579        let mut layer =
580            GoogleAuthMiddlewareLayer::new(token_generator, Some("projects/first".to_string()))
581                .unwrap()
582                .amend_user_agent("extra-ua".to_string())
583                .unwrap();
584        layer.set_additional_headers(extra);
585        let mut first = layer.layer(DummyService {
586            tx: Arc::new(tx.clone()),
587        });
588
589        let (other_tx, mut other_rx) = tokio::sync::mpsc::channel(2);
590        let other = DummyService {
591            tx: Arc::new(other_tx),
592        };
593        let mut with_prefix = first
594            .with_inner(other.clone(), Some("projects/second".to_string()))
595            .unwrap();
596        let mut without_prefix = first.with_inner(other, None).unwrap();
597
598        for service in [&mut first, &mut with_prefix, &mut without_prefix] {
599            let req = Request::builder()
600                .uri("http://example.com")
601                .body("".to_string())
602                .unwrap();
603            tower::Service::call(service, req).await.unwrap();
604        }
605
606        let first_req = rx.recv().await.unwrap();
607        let with_prefix_req = other_rx.recv().await.unwrap();
608        let without_prefix_req = other_rx.recv().await.unwrap();
609
610        let expected_ua = format!("gcloud-sdk-rs/{} extra-ua", env!("CARGO_PKG_VERSION"));
611        for req in [&first_req, &with_prefix_req, &without_prefix_req] {
612            assert_eq!(
613                req.headers().get(hyper::header::USER_AGENT).unwrap(),
614                expected_ua.as_str()
615            );
616            assert_eq!(req.headers().get("x-test-header").unwrap(), "test-value");
617            assert_eq!(
618                req.headers().get("authorization").unwrap(),
619                "Bearer counted-token"
620            );
621        }
622        assert_eq!(
623            first_req
624                .headers()
625                .get(GOOGLE_CLOUD_RESOURCE_PREFIX)
626                .unwrap(),
627            "projects/first"
628        );
629        assert_eq!(
630            with_prefix_req
631                .headers()
632                .get(GOOGLE_CLOUD_RESOURCE_PREFIX)
633                .unwrap(),
634            "projects/second"
635        );
636        assert!(without_prefix_req
637            .headers()
638            .get(GOOGLE_CLOUD_RESOURCE_PREFIX)
639            .is_none());
640        assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 1);
641    }
642}
gcloud_sdk/middleware.rs

gcloud_sdk/
middleware.rs
