Skip to main content

gcloud_sdk/rest_apis/
rest_api_client.rs

1use crate::{GoogleAuthTokenGenerator, TokenSourceType, GCP_DEFAULT_SCOPES};
2use async_trait::async_trait;
3use hyper::Uri;
4use reqwest::{IntoUrl, Method, Request, RequestBuilder};
5use std::sync::Arc;
6
7#[derive(Clone)]
8pub struct GoogleRestApi {
9    pub client: reqwest::Client,
10    pub token_generator: Arc<GoogleAuthTokenGenerator>,
11}
12
13impl GoogleRestApi {
14    pub async fn new() -> crate::error::Result<Self> {
15        Self::with_token_source(TokenSourceType::Default, GCP_DEFAULT_SCOPES.clone()).await
16    }
17
18    pub async fn with_token_source(
19        token_source_type: TokenSourceType,
20        token_scopes: Vec<String>,
21    ) -> crate::error::Result<Self> {
22        let client = reqwest::Client::new();
23        Self::with_client_token_source(client, token_source_type, token_scopes).await
24    }
25
26    pub async fn with_client_token_source(
27        client: reqwest::Client,
28        token_source_type: TokenSourceType,
29        token_scopes: Vec<String>,
30    ) -> crate::error::Result<Self> {
31        let token_generator =
32            GoogleAuthTokenGenerator::new(token_source_type, token_scopes).await?;
33
34        Ok(Self {
35            client,
36            token_generator: Arc::new(token_generator),
37        })
38    }
39
40    pub async fn with_google_token<'a>(
41        &self,
42        request: RequestBuilder,
43    ) -> crate::error::Result<RequestBuilder> {
44        let authorization = self.token_generator.authorization_header().await?;
45        Ok(request.header(reqwest::header::AUTHORIZATION, authorization))
46    }
47
48    /// A [`GoogleAuthReqwestMiddleware`](crate::GoogleAuthReqwestMiddleware) over this
49    /// client's token generator, so requests sent through a `reqwest-middleware` client
50    /// reuse the tokens this client has already minted.
51    #[cfg(feature = "reqwest-middleware")]
52    pub fn middleware(&self) -> crate::GoogleAuthReqwestMiddleware {
53        crate::GoogleAuthReqwestMiddleware::new(self.token_generator.clone())
54    }
55
56    pub async fn get<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
57        self.with_google_token(self.client.request(Method::GET, url))
58            .await
59    }
60    pub async fn post<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
61        self.with_google_token(self.client.request(Method::POST, url))
62            .await
63    }
64    pub async fn put<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
65        self.with_google_token(self.client.request(Method::PUT, url))
66            .await
67    }
68    pub async fn patch<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
69        self.with_google_token(self.client.request(Method::PATCH, url))
70            .await
71    }
72    pub async fn delete<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
73        self.with_google_token(self.client.request(Method::DELETE, url))
74            .await
75    }
76    pub async fn head<U: IntoUrl>(&self, url: U) -> crate::error::Result<RequestBuilder> {
77        self.with_google_token(self.client.request(Method::HEAD, url))
78            .await
79    }
80}
81
82pub fn create_hyper_uri_with_params<'p, PT, TS>(url_str: &str, params: &'p PT) -> Uri
83where
84    PT: std::iter::IntoIterator<Item = (&'p str, Option<&'p TS>)> + Clone,
85    TS: std::string::ToString + 'p,
86{
87    let url_query_params: Vec<(String, String)> = params
88        .clone()
89        .into_iter()
90        .filter_map(|(k, vo)| vo.map(|v| (k.to_string(), v.to_string())))
91        .collect();
92
93    let url: url::Url = url::Url::parse_with_params(url_str, url_query_params)
94        .unwrap()
95        .as_str()
96        .parse()
97        .unwrap();
98
99    url.as_str().parse().unwrap()
100}
101
102#[cfg(all(test, feature = "reqwest-middleware"))]
103mod tests {
104    use super::*;
105    use crate::test_support::{StubResponse, StubServer};
106    use crate::token_source::{Source, Token};
107    use jiff::{SignedDuration, Timestamp};
108    use secret_vault_value::SecretValue;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110
111    struct CountingSource {
112        mints: Arc<AtomicUsize>,
113    }
114
115    #[async_trait]
116    impl Source for CountingSource {
117        async fn token(&self) -> crate::error::Result<Token> {
118            self.mints.fetch_add(1, Ordering::SeqCst);
119            Ok(Token::new(
120                "Bearer".to_string(),
121                SecretValue::from("shared-token"),
122                Timestamp::now() + SignedDuration::from_hours(1),
123            ))
124        }
125    }
126
127    #[tokio::test]
128    async fn middleware_shares_the_token_cache() {
129        let mints = Arc::new(AtomicUsize::new(0));
130        let api = GoogleRestApi {
131            client: reqwest::Client::new(),
132            token_generator: Arc::new(GoogleAuthTokenGenerator::from_source(Box::new(
133                CountingSource {
134                    mints: mints.clone(),
135                },
136            )
137                as crate::BoxSource)),
138        };
139        let service = StubServer::start(vec![
140            StubResponse::json("200 OK", "{}"),
141            StubResponse::json("200 OK", "{}"),
142        ])
143        .await;
144        let middleware_client = reqwest_middleware::ClientBuilder::new(reqwest::Client::new())
145            .with(api.middleware())
146            .build();
147
148        api.get(format!("{}/rest", service.url))
149            .await
150            .unwrap()
151            .send()
152            .await
153            .unwrap();
154        middleware_client
155            .get(format!("{}/middleware", service.url))
156            .send()
157            .await
158            .unwrap();
159
160        let received = service.received();
161        assert_eq!(
162            received[0].header("authorization"),
163            Some("Bearer shared-token")
164        );
165        assert_eq!(
166            received[1].header("authorization"),
167            Some("Bearer shared-token")
168        );
169        assert_eq!(mints.load(Ordering::SeqCst), 1);
170    }
171}