gcloud_sdk/
middleware.rs

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

gcloud_sdk/
middleware.rs
