olai_http/gcp/
builder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::str::FromStr;
19use std::sync::Arc;
20use std::time::Duration;
21
22use serde::{Deserialize, Serialize};
23use tokio::runtime::Handle;
24
25use crate::gcp::GoogleConfig;
26use crate::gcp::credential::GcpCredential;
27use crate::gcp::credential::{
28    ApplicationDefaultCredentials, GcpWorkloadIdentityProvider, InstanceCredentialProvider,
29    ServiceAccountCredentials,
30};
31use crate::gcp::{GcpCredentialProvider, credential};
32use crate::service::make_service;
33use crate::{
34    ClientConfigKey, ClientOptions, Result, RetryConfig, StaticCredentialProvider,
35    TokenCredentialProvider,
36};
37
38const TOKEN_MIN_TTL: Duration = Duration::from_secs(4 * 60);
39
40#[derive(Debug, thiserror::Error)]
41enum Error {
42    #[error("One of service account path or service account key may be provided.")]
43    ServiceAccountPathAndKeyProvided,
44
45    #[error("Configuration key: '{}' is not known.", key)]
46    UnknownConfigurationKey { key: String },
47
48    #[error("GCP credential error: {}", source)]
49    Credential { source: credential::Error },
50}
51
52impl From<Error> for crate::Error {
53    fn from(err: Error) -> Self {
54        match err {
55            Error::UnknownConfigurationKey { key } => Self::UnknownConfigurationKey { key },
56            _ => Self::Generic {
57                source: Box::new(err),
58            },
59        }
60    }
61}
62
63/// Configure a connection to Google Cloud Storage.
64///
65/// If no credentials are explicitly provided, they will be sourced
66/// from the environment as documented [here](https://cloud.google.com/docs/authentication/application-default-credentials).
67///
68/// # Example
69/// ```
70/// # use olai_http::gcp::GoogleBuilder;
71/// let gcs = GoogleBuilder::from_env().build(None);
72/// ```
73#[derive(Debug, Clone)]
74pub struct GoogleBuilder {
75    /// Path to the service account file
76    service_account_path: Option<String>,
77    /// The serialized service account key
78    service_account_key: Option<String>,
79    /// Path to the application credentials file.
80    application_credentials_path: Option<String>,
81    /// Retry config
82    retry_config: RetryConfig,
83    /// Client options
84    client_options: ClientOptions,
85    /// Credentials
86    credentials: Option<GcpCredentialProvider>,
87}
88
89/// Configuration keys for [`GoogleBuilder`]
90///
91/// Configuration via keys can be done via [`GoogleBuilder::with_config`]
92///
93/// # Example
94/// ```
95/// # use olai_http::gcp::{GoogleBuilder, GoogleConfigKey};
96/// let builder = GoogleBuilder::new()
97///     .with_config("google_service_account".parse().unwrap(), "my-service-account");
98/// ```
99#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Serialize, Deserialize)]
100#[non_exhaustive]
101pub enum GoogleConfigKey {
102    /// Path to the service account file
103    ///
104    /// Supported keys:
105    /// - `google_service_account`
106    /// - `service_account`
107    /// - `google_service_account_path`
108    /// - `service_account_path`
109    ServiceAccount,
110
111    /// The serialized service account key.
112    ///
113    /// Supported keys:
114    /// - `google_service_account_key`
115    /// - `service_account_key`
116    ServiceAccountKey,
117
118    /// Application credentials path
119    ///
120    /// See [`GoogleBuilder::with_application_credentials`].
121    ApplicationCredentials,
122
123    /// Client options
124    Client(ClientConfigKey),
125}
126
127impl AsRef<str> for GoogleConfigKey {
128    fn as_ref(&self) -> &str {
129        match self {
130            Self::ServiceAccount => "google_service_account",
131            Self::ServiceAccountKey => "google_service_account_key",
132            Self::ApplicationCredentials => "google_application_credentials",
133            Self::Client(key) => key.as_ref(),
134        }
135    }
136}
137
138impl FromStr for GoogleConfigKey {
139    type Err = crate::Error;
140
141    fn from_str(s: &str) -> Result<Self, Self::Err> {
142        match s {
143            "google_service_account"
144            | "service_account"
145            | "google_service_account_path"
146            | "service_account_path" => Ok(Self::ServiceAccount),
147            "google_service_account_key" | "service_account_key" => Ok(Self::ServiceAccountKey),
148            "google_application_credentials" => Ok(Self::ApplicationCredentials),
149            _ => match s.strip_prefix("google_").unwrap_or(s).parse() {
150                Ok(key) => Ok(Self::Client(key)),
151                Err(_) => Err(Error::UnknownConfigurationKey { key: s.into() }.into()),
152            },
153        }
154    }
155}
156
157impl Default for GoogleBuilder {
158    fn default() -> Self {
159        Self {
160            service_account_path: None,
161            service_account_key: None,
162            application_credentials_path: None,
163            retry_config: Default::default(),
164            client_options: ClientOptions::new().with_allow_http(true),
165            credentials: None,
166        }
167    }
168}
169
170impl GoogleBuilder {
171    /// Create a new [`GoogleBuilder`] with default values.
172    pub fn new() -> Self {
173        Default::default()
174    }
175
176    /// Create an instance of [`GoogleBuilder`] with values pre-populated from environment variables.
177    ///
178    /// Variables extracted from environment:
179    /// * GOOGLE_SERVICE_ACCOUNT: location of service account file
180    /// * GOOGLE_SERVICE_ACCOUNT_PATH: (alias) location of service account file
181    /// * SERVICE_ACCOUNT: (alias) location of service account file
182    /// * GOOGLE_SERVICE_ACCOUNT_KEY: JSON serialized service account key
183    ///
184    /// # Example
185    /// ```
186    /// use olai_http::gcp::GoogleBuilder;
187    ///
188    /// let gcs = GoogleBuilder::from_env()
189    ///     .build(None);
190    /// ```
191    pub fn from_env() -> Self {
192        let mut builder = Self::default();
193
194        if let Ok(service_account_path) = std::env::var("SERVICE_ACCOUNT") {
195            builder.service_account_path = Some(service_account_path);
196        }
197
198        for (os_key, os_value) in std::env::vars_os() {
199            if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str())
200                && key.starts_with("GOOGLE_")
201                && let Ok(config_key) = key.to_ascii_lowercase().parse()
202            {
203                builder = builder.with_config(config_key, value);
204            }
205        }
206
207        builder
208    }
209
210    /// Set an option on the builder via a key - value pair.
211    pub fn with_config(mut self, key: GoogleConfigKey, value: impl Into<String>) -> Self {
212        match key {
213            GoogleConfigKey::ServiceAccount => self.service_account_path = Some(value.into()),
214            GoogleConfigKey::ServiceAccountKey => self.service_account_key = Some(value.into()),
215            GoogleConfigKey::ApplicationCredentials => {
216                self.application_credentials_path = Some(value.into())
217            }
218            GoogleConfigKey::Client(key) => {
219                self.client_options = self.client_options.with_config(key, value)
220            }
221        };
222        self
223    }
224
225    /// Get config value via a [`GoogleConfigKey`].
226    ///
227    /// # Example
228    /// ```
229    /// use olai_http::gcp::{GoogleBuilder, GoogleConfigKey};
230    ///
231    /// let builder = GoogleBuilder::from_env()
232    ///     .with_service_account_key("foo");
233    /// let service_account_key = builder.get_config_value(&GoogleConfigKey::ServiceAccountKey).unwrap_or_default();
234    /// assert_eq!("foo", &service_account_key);
235    /// ```
236    pub fn get_config_value(&self, key: &GoogleConfigKey) -> Option<String> {
237        match key {
238            GoogleConfigKey::ServiceAccount => self.service_account_path.clone(),
239            GoogleConfigKey::ServiceAccountKey => self.service_account_key.clone(),
240            GoogleConfigKey::ApplicationCredentials => self.application_credentials_path.clone(),
241            GoogleConfigKey::Client(key) => self.client_options.get_config_value(key),
242        }
243    }
244
245    /// Set the path to the service account file.
246    ///
247    /// This or [`GoogleBuilder::with_service_account_key`] must be
248    /// set.
249    ///
250    /// Example `"/tmp/gcs.json"`.
251    ///
252    /// Example contents of `gcs.json`:
253    ///
254    /// ```json
255    /// {
256    ///    "gcs_base_url": "https://localhost:4443",
257    ///    "disable_oauth": true,
258    ///    "client_email": "",
259    ///    "private_key": ""
260    /// }
261    /// ```
262    pub fn with_service_account_path(mut self, service_account_path: impl Into<String>) -> Self {
263        self.service_account_path = Some(service_account_path.into());
264        self
265    }
266
267    /// Set the service account key. The service account must be in the JSON
268    /// format.
269    ///
270    /// This or [`GoogleBuilder::with_service_account_path`] must be
271    /// set.
272    pub fn with_service_account_key(mut self, service_account: impl Into<String>) -> Self {
273        self.service_account_key = Some(service_account.into());
274        self
275    }
276
277    /// Set the path to the application credentials file.
278    ///
279    /// <https://cloud.google.com/docs/authentication/provide-credentials-adc>
280    pub fn with_application_credentials(
281        mut self,
282        application_credentials_path: impl Into<String>,
283    ) -> Self {
284        self.application_credentials_path = Some(application_credentials_path.into());
285        self
286    }
287
288    /// Set the credential provider overriding any other options
289    pub fn with_credentials(mut self, credentials: GcpCredentialProvider) -> Self {
290        self.credentials = Some(credentials);
291        self
292    }
293
294    /// Set the retry configuration
295    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
296        self.retry_config = retry_config;
297        self
298    }
299
300    /// Set the proxy_url to be used by the underlying client
301    pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
302        self.client_options = self.client_options.with_proxy_url(proxy_url);
303        self
304    }
305
306    /// Set a trusted proxy CA certificate
307    pub fn with_proxy_ca_certificate(mut self, proxy_ca_certificate: impl Into<String>) -> Self {
308        self.client_options = self
309            .client_options
310            .with_proxy_ca_certificate(proxy_ca_certificate);
311        self
312    }
313
314    /// Set a list of hosts to exclude from proxy connections
315    pub fn with_proxy_excludes(mut self, proxy_excludes: impl Into<String>) -> Self {
316        self.client_options = self.client_options.with_proxy_excludes(proxy_excludes);
317        self
318    }
319
320    /// Sets the client options, overriding any already set
321    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
322        self.client_options = options;
323        self
324    }
325
326    /// Configure a connection to Google Cloud Storage, returning a
327    /// new [`GoogleConfig`] and consuming `self`
328    ///
329    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
330    /// will be spawned on the given runtime handle.
331    pub fn build(self, runtime: Option<&Handle>) -> Result<GoogleConfig> {
332        // First try to initialize from the service account information.
333        let service_account_credentials =
334            match (self.service_account_path, self.service_account_key) {
335                (Some(path), None) => Some(
336                    ServiceAccountCredentials::from_file(path)
337                        .map_err(|source| Error::Credential { source })?,
338                ),
339                (None, Some(key)) => Some(
340                    ServiceAccountCredentials::from_key(&key)
341                        .map_err(|source| Error::Credential { source })?,
342                ),
343                (None, None) => None,
344                (Some(_), Some(_)) => return Err(Error::ServiceAccountPathAndKeyProvided.into()),
345            };
346
347        // Then try to initialize from the application credentials file, or the environment.
348        let application_default_credentials =
349            ApplicationDefaultCredentials::read(self.application_credentials_path.as_deref())?;
350
351        let disable_oauth = service_account_credentials
352            .as_ref()
353            .map(|c| c.disable_oauth)
354            .unwrap_or(false);
355
356        let credentials = if let Some(credentials) = self.credentials {
357            credentials
358        } else if disable_oauth {
359            Arc::new(StaticCredentialProvider::new(GcpCredential {
360                bearer: "".to_string(),
361            })) as _
362        } else if let Some(credentials) = service_account_credentials.clone() {
363            let client = self.client_options.client()?;
364            let service = make_service(client.clone(), runtime);
365            Arc::new(TokenCredentialProvider::new(
366                credentials.token_provider()?,
367                client,
368                service,
369                self.retry_config.clone(),
370            )) as _
371        } else if let Some(credentials) = application_default_credentials.clone() {
372            match credentials {
373                ApplicationDefaultCredentials::AuthorizedUser(token) => {
374                    let client = self.client_options.client()?;
375                    let service = make_service(client.clone(), runtime);
376                    Arc::new(
377                        TokenCredentialProvider::new(
378                            token,
379                            client,
380                            service,
381                            self.retry_config.clone(),
382                        )
383                        .with_min_ttl(TOKEN_MIN_TTL),
384                    ) as _
385                }
386                ApplicationDefaultCredentials::ServiceAccount(token) => {
387                    let client = self.client_options.client()?;
388                    let service = make_service(client.clone(), runtime);
389                    Arc::new(TokenCredentialProvider::new(
390                        token.token_provider()?,
391                        client,
392                        service,
393                        self.retry_config.clone(),
394                    )) as _
395                }
396                ApplicationDefaultCredentials::ExternalAccount(ext) => {
397                    let client = self.client_options.client()?;
398                    let service = make_service(client.clone(), runtime);
399                    let sts_endpoint = ext.token_url.clone();
400                    Arc::new(
401                        TokenCredentialProvider::new(
402                            GcpWorkloadIdentityProvider {
403                                credentials: ext,
404                                sts_endpoint,
405                            },
406                            client,
407                            service,
408                            self.retry_config.clone(),
409                        )
410                        .with_min_ttl(TOKEN_MIN_TTL),
411                    ) as _
412                }
413            }
414        } else {
415            let client = self.client_options.metadata_client()?;
416            let service = make_service(client.clone(), runtime);
417            Arc::new(
418                TokenCredentialProvider::new(
419                    InstanceCredentialProvider::default(),
420                    client,
421                    service,
422                    self.retry_config.clone(),
423                )
424                .with_min_ttl(TOKEN_MIN_TTL),
425            ) as _
426        };
427
428        Ok(GoogleConfig::new(
429            credentials,
430            self.retry_config,
431            self.client_options,
432        ))
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    use std::collections::HashMap;
440    use std::io::Write;
441    use tempfile::NamedTempFile;
442
443    const FAKE_KEY: &str = r#"{"private_key": "private_key", "private_key_id": "private_key_id", "client_email":"client_email", "disable_oauth":true}"#;
444
445    #[test]
446    fn gcs_test_service_account_key_and_path() {
447        let mut tfile = NamedTempFile::new().unwrap();
448        write!(tfile, "{FAKE_KEY}").unwrap();
449        let _ = GoogleBuilder::new()
450            .with_service_account_key(FAKE_KEY)
451            .with_service_account_path(tfile.path().to_str().unwrap())
452            .build(None)
453            .unwrap_err();
454    }
455
456    #[test]
457    fn gcs_test_config_from_map() {
458        let google_service_account = "object_store:fake_service_account".to_string();
459        let options = HashMap::from([("google_service_account", google_service_account.clone())]);
460
461        let builder = options
462            .iter()
463            .fold(GoogleBuilder::new(), |builder, (key, value)| {
464                builder.with_config(key.parse().unwrap(), value)
465            });
466
467        assert_eq!(
468            builder.service_account_path.unwrap(),
469            google_service_account.as_str()
470        );
471    }
472
473    #[test]
474    fn gcs_test_config_aliases() {
475        // Service account path
476        for alias in [
477            "google_service_account",
478            "service_account",
479            "google_service_account_path",
480            "service_account_path",
481        ] {
482            let builder =
483                GoogleBuilder::new().with_config(alias.parse().unwrap(), "/fake/path.json");
484            assert_eq!("/fake/path.json", builder.service_account_path.unwrap());
485        }
486
487        // Service account key
488        for alias in ["google_service_account_key", "service_account_key"] {
489            let builder = GoogleBuilder::new().with_config(alias.parse().unwrap(), FAKE_KEY);
490            assert_eq!(FAKE_KEY, builder.service_account_key.unwrap());
491        }
492    }
493
494    #[tokio::test]
495    async fn gcs_test_proxy_url() {
496        let mut tfile = NamedTempFile::new().unwrap();
497        write!(tfile, "{FAKE_KEY}").unwrap();
498        let service_account_path = tfile.path();
499        let gcs = GoogleBuilder::new()
500            .with_service_account_path(service_account_path.to_str().unwrap())
501            .with_proxy_url("https://example.com")
502            .build(None);
503        assert!(gcs.is_ok());
504    }
505
506    #[test]
507    fn gcs_test_service_account_key_only() {
508        let _ = GoogleBuilder::new()
509            .with_service_account_key(FAKE_KEY)
510            .build(None)
511            .unwrap();
512    }
513
514    #[test]
515    fn gcs_test_config_get_value() {
516        let google_service_account = "object_store:fake_service_account".to_string();
517        let builder = GoogleBuilder::new()
518            .with_config(GoogleConfigKey::ServiceAccount, &google_service_account);
519
520        assert_eq!(
521            builder
522                .get_config_value(&GoogleConfigKey::ServiceAccount)
523                .unwrap(),
524            google_service_account
525        );
526    }
527
528    #[test]
529    fn gcp_test_client_opts() {
530        let key = "GOOGLE_PROXY_URL";
531        if let Ok(config_key) = key.to_ascii_lowercase().parse() {
532            assert_eq!(
533                GoogleConfigKey::Client(ClientConfigKey::ProxyUrl),
534                config_key
535            );
536        } else {
537            panic!("{key} not propagated as ClientConfigKey");
538        }
539    }
540}
olai_http/gcp/builder.rs

olai_http/gcp/
builder.rs
