gcloud_sdk/
middleware.rs

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

gcloud_sdk/
middleware.rs
