olai_http/azure/
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;
20
21use serde::{Deserialize, Serialize};
22use tokio::runtime::Handle;
23
24use super::AzureConfig;
25use crate::azure::credential::{
26    AzureCliCredential, ClientSecretOAuthProvider, ImdsManagedIdentityProvider,
27    WorkloadIdentityOAuthProvider,
28};
29use crate::azure::{AzureCredential, AzureCredentialProvider};
30use crate::config::ConfigValue;
31use crate::service::make_service;
32use crate::{
33    ClientConfigKey, ClientOptions, Result, RetryConfig, StaticCredentialProvider,
34    TokenCredentialProvider,
35};
36
37const MSI_ENDPOINT_ENV_KEY: &str = "IDENTITY_ENDPOINT";
38
39#[derive(Debug, thiserror::Error)]
40enum Error {
41    #[error("OAuth scope must be set when using client secret authentication")]
42    MissingScope,
43
44    #[error("Configuration key: '{}' is not known.", key)]
45    UnknownConfigurationKey { key: String },
46}
47
48impl From<Error> for crate::Error {
49    fn from(source: Error) -> Self {
50        match source {
51            Error::UnknownConfigurationKey { key } => Self::UnknownConfigurationKey { key },
52            Error::MissingScope => Self::Generic {
53                source: Box::new(source),
54            },
55        }
56    }
57}
58
59/// Configure Azure authentication credentials.
60///
61/// # Example
62/// ```
63/// # let ACCOUNT = "foo";
64/// # let ACCESS_KEY = "foo";
65/// # use olai_http::azure::AzureBuilder;
66/// let config = AzureBuilder::new()
67///  .with_account(ACCOUNT)
68///  .with_access_key(ACCESS_KEY)
69///  .build(None);
70/// ```
71#[derive(Default, Clone)]
72pub struct AzureBuilder {
73    account_name: Option<String>,
74    access_key: Option<String>,
75    bearer_token: Option<String>,
76    client_id: Option<String>,
77    client_secret: Option<String>,
78    tenant_id: Option<String>,
79    scope: Option<String>,
80    authority_host: Option<String>,
81    endpoint: Option<String>,
82    msi_endpoint: Option<String>,
83    object_id: Option<String>,
84    msi_resource_id: Option<String>,
85    federated_token_file: Option<String>,
86    use_azure_cli: ConfigValue<bool>,
87    retry_config: RetryConfig,
88    client_options: ClientOptions,
89    credentials: Option<AzureCredentialProvider>,
90    skip_signature: ConfigValue<bool>,
91}
92
93/// Configuration keys for [`AzureBuilder`]
94///
95/// Configuration via keys can be done via [`AzureBuilder::with_config`]
96///
97/// # Example
98/// ```
99/// # use olai_http::azure::{AzureBuilder, AzureConfigKey};
100/// let builder = AzureBuilder::new()
101///     .with_config("azure_client_id".parse().unwrap(), "my-client-id")
102///     .with_config(AzureConfigKey::AuthorityId, "my-tenant-id");
103/// ```
104#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Deserialize, Serialize)]
105#[non_exhaustive]
106pub enum AzureConfigKey {
107    /// The name of the azure storage account
108    ///
109    /// Supported keys:
110    /// - `azure_storage_account_name`
111    /// - `account_name`
112    AccountName,
113
114    /// Master key for accessing storage account
115    ///
116    /// Supported keys:
117    /// - `azure_storage_account_key`
118    /// - `azure_storage_access_key`
119    /// - `azure_storage_master_key`
120    /// - `access_key`
121    /// - `account_key`
122    /// - `master_key`
123    AccessKey,
124
125    /// Service principal client id for authorizing requests
126    ///
127    /// Supported keys:
128    /// - `azure_storage_client_id`
129    /// - `azure_client_id`
130    /// - `client_id`
131    ClientId,
132
133    /// Service principal client secret for authorizing requests
134    ///
135    /// Supported keys:
136    /// - `azure_storage_client_secret`
137    /// - `azure_client_secret`
138    /// - `client_secret`
139    ClientSecret,
140
141    /// Tenant id used in oauth flows
142    ///
143    /// Supported keys:
144    /// - `azure_storage_tenant_id`
145    /// - `azure_storage_authority_id`
146    /// - `azure_tenant_id`
147    /// - `azure_authority_id`
148    /// - `tenant_id`
149    /// - `authority_id`
150    AuthorityId,
151
152    /// Authority host used in oauth flows
153    ///
154    /// Supported keys:
155    /// - `azure_storage_authority_host`
156    /// - `azure_authority_host`
157    /// - `authority_host`
158    AuthorityHost,
159
160    /// OAuth scope for client credentials flow
161    ///
162    /// Supported keys:
163    /// - `azure_scope`
164    /// - `scope`
165    Scope,
166
167    /// Bearer token
168    ///
169    /// Supported keys:
170    /// - `azure_storage_token`
171    /// - `bearer_token`
172    /// - `token`
173    Token,
174
175    /// Override the endpoint used to communicate with Azure
176    ///
177    /// Supported keys:
178    /// - `azure_storage_endpoint`
179    /// - `azure_endpoint`
180    /// - `endpoint`
181    Endpoint,
182
183    /// Endpoint to request a imds managed identity token
184    ///
185    /// Supported keys:
186    /// - `azure_msi_endpoint`
187    /// - `azure_identity_endpoint`
188    /// - `identity_endpoint`
189    /// - `msi_endpoint`
190    MsiEndpoint,
191
192    /// Object id for use with managed identity authentication
193    ///
194    /// Supported keys:
195    /// - `azure_object_id`
196    /// - `object_id`
197    ObjectId,
198
199    /// Msi resource id for use with managed identity authentication
200    ///
201    /// Supported keys:
202    /// - `azure_msi_resource_id`
203    /// - `msi_resource_id`
204    MsiResourceId,
205
206    /// File containing token for Azure AD workload identity federation
207    ///
208    /// Supported keys:
209    /// - `azure_federated_token_file`
210    /// - `federated_token_file`
211    FederatedTokenFile,
212
213    /// Use azure cli for acquiring access token
214    ///
215    /// Supported keys:
216    /// - `azure_use_azure_cli`
217    /// - `use_azure_cli`
218    UseAzureCli,
219
220    /// Skip signing requests
221    ///
222    /// Supported keys:
223    /// - `azure_skip_signature`
224    /// - `skip_signature`
225    SkipSignature,
226
227    /// Client options
228    Client(ClientConfigKey),
229}
230
231impl AsRef<str> for AzureConfigKey {
232    fn as_ref(&self) -> &str {
233        match self {
234            Self::AccountName => "azure_storage_account_name",
235            Self::AccessKey => "azure_storage_account_key",
236            Self::ClientId => "azure_storage_client_id",
237            Self::ClientSecret => "azure_storage_client_secret",
238            Self::AuthorityId => "azure_storage_tenant_id",
239            Self::AuthorityHost => "azure_storage_authority_host",
240            Self::Scope => "azure_scope",
241            Self::Token => "azure_storage_token",
242            Self::Endpoint => "azure_storage_endpoint",
243            Self::MsiEndpoint => "azure_msi_endpoint",
244            Self::ObjectId => "azure_object_id",
245            Self::MsiResourceId => "azure_msi_resource_id",
246            Self::FederatedTokenFile => "azure_federated_token_file",
247            Self::UseAzureCli => "azure_use_azure_cli",
248            Self::SkipSignature => "azure_skip_signature",
249            Self::Client(key) => key.as_ref(),
250        }
251    }
252}
253
254impl FromStr for AzureConfigKey {
255    type Err = crate::Error;
256
257    fn from_str(s: &str) -> Result<Self, Self::Err> {
258        match s {
259            "azure_storage_account_key"
260            | "azure_storage_access_key"
261            | "azure_storage_master_key"
262            | "master_key"
263            | "account_key"
264            | "access_key" => Ok(Self::AccessKey),
265            "azure_storage_account_name" | "account_name" => Ok(Self::AccountName),
266            "azure_storage_client_id" | "azure_client_id" | "client_id" => Ok(Self::ClientId),
267            "azure_storage_client_secret" | "azure_client_secret" | "client_secret" => {
268                Ok(Self::ClientSecret)
269            }
270            "azure_storage_tenant_id"
271            | "azure_storage_authority_id"
272            | "azure_tenant_id"
273            | "azure_authority_id"
274            | "tenant_id"
275            | "authority_id" => Ok(Self::AuthorityId),
276            "azure_storage_authority_host" | "azure_authority_host" | "authority_host" => {
277                Ok(Self::AuthorityHost)
278            }
279            "azure_scope" | "scope" => Ok(Self::Scope),
280            "azure_storage_token" | "bearer_token" | "token" => Ok(Self::Token),
281            "azure_storage_endpoint" | "azure_endpoint" | "endpoint" => Ok(Self::Endpoint),
282            "azure_msi_endpoint"
283            | "azure_identity_endpoint"
284            | "identity_endpoint"
285            | "msi_endpoint" => Ok(Self::MsiEndpoint),
286            "azure_object_id" | "object_id" => Ok(Self::ObjectId),
287            "azure_msi_resource_id" | "msi_resource_id" => Ok(Self::MsiResourceId),
288            "azure_federated_token_file" | "federated_token_file" => Ok(Self::FederatedTokenFile),
289            "azure_use_azure_cli" | "use_azure_cli" => Ok(Self::UseAzureCli),
290            "azure_skip_signature" | "skip_signature" => Ok(Self::SkipSignature),
291            "azure_allow_http" => Ok(Self::Client(ClientConfigKey::AllowHttp)),
292            _ => match s.strip_prefix("azure_").unwrap_or(s).parse() {
293                Ok(key) => Ok(Self::Client(key)),
294                Err(_) => Err(Error::UnknownConfigurationKey { key: s.into() }.into()),
295            },
296        }
297    }
298}
299
300impl std::fmt::Debug for AzureBuilder {
301    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
302        write!(f, "AzureBuilder {{ account: {:?} }}", self.account_name)
303    }
304}
305
306impl AzureBuilder {
307    /// Create a new [`AzureBuilder`] with default values.
308    pub fn new() -> Self {
309        Default::default()
310    }
311
312    /// Create an instance of [`AzureBuilder`] with values pre-populated from environment variables.
313    ///
314    /// Variables extracted from environment:
315    /// * AZURE_STORAGE_ACCOUNT_NAME: storage account name
316    /// * AZURE_STORAGE_ACCOUNT_KEY: storage account master key
317    /// * AZURE_STORAGE_ACCESS_KEY: alias for AZURE_STORAGE_ACCOUNT_KEY
318    /// * AZURE_STORAGE_CLIENT_ID -> client id for service principal authorization
319    /// * AZURE_STORAGE_CLIENT_SECRET -> client secret for service principal authorization
320    /// * AZURE_STORAGE_TENANT_ID -> tenant id used in oauth flows
321    pub fn from_env() -> Self {
322        let mut builder = Self::default();
323        for (os_key, os_value) in std::env::vars_os() {
324            if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) {
325                if key.starts_with("AZURE_") {
326                    if let Ok(config_key) = key.to_ascii_lowercase().parse() {
327                        builder = builder.with_config(config_key, value);
328                    }
329                }
330            }
331        }
332
333        if let Ok(text) = std::env::var(MSI_ENDPOINT_ENV_KEY) {
334            builder = builder.with_msi_endpoint(text);
335        }
336
337        builder
338    }
339
340    /// Set an option on the builder via a key - value pair.
341    pub fn with_config(mut self, key: AzureConfigKey, value: impl Into<String>) -> Self {
342        match key {
343            AzureConfigKey::AccessKey => self.access_key = Some(value.into()),
344            AzureConfigKey::AccountName => self.account_name = Some(value.into()),
345            AzureConfigKey::ClientId => self.client_id = Some(value.into()),
346            AzureConfigKey::ClientSecret => self.client_secret = Some(value.into()),
347            AzureConfigKey::AuthorityId => self.tenant_id = Some(value.into()),
348            AzureConfigKey::AuthorityHost => self.authority_host = Some(value.into()),
349            AzureConfigKey::Scope => self.scope = Some(value.into()),
350            AzureConfigKey::Token => self.bearer_token = Some(value.into()),
351            AzureConfigKey::MsiEndpoint => self.msi_endpoint = Some(value.into()),
352            AzureConfigKey::ObjectId => self.object_id = Some(value.into()),
353            AzureConfigKey::MsiResourceId => self.msi_resource_id = Some(value.into()),
354            AzureConfigKey::FederatedTokenFile => self.federated_token_file = Some(value.into()),
355            AzureConfigKey::UseAzureCli => self.use_azure_cli.parse(value),
356            AzureConfigKey::SkipSignature => self.skip_signature.parse(value),
357            AzureConfigKey::Endpoint => self.endpoint = Some(value.into()),
358            AzureConfigKey::Client(key) => {
359                self.client_options = self.client_options.with_config(key, value)
360            }
361        };
362        self
363    }
364
365    /// Get config value via a [`AzureConfigKey`].
366    pub fn get_config_value(&self, key: &AzureConfigKey) -> Option<String> {
367        match key {
368            AzureConfigKey::AccountName => self.account_name.clone(),
369            AzureConfigKey::AccessKey => self.access_key.clone(),
370            AzureConfigKey::ClientId => self.client_id.clone(),
371            AzureConfigKey::ClientSecret => self.client_secret.clone(),
372            AzureConfigKey::AuthorityId => self.tenant_id.clone(),
373            AzureConfigKey::AuthorityHost => self.authority_host.clone(),
374            AzureConfigKey::Scope => self.scope.clone(),
375            AzureConfigKey::Token => self.bearer_token.clone(),
376            AzureConfigKey::Endpoint => self.endpoint.clone(),
377            AzureConfigKey::MsiEndpoint => self.msi_endpoint.clone(),
378            AzureConfigKey::ObjectId => self.object_id.clone(),
379            AzureConfigKey::MsiResourceId => self.msi_resource_id.clone(),
380            AzureConfigKey::FederatedTokenFile => self.federated_token_file.clone(),
381            AzureConfigKey::UseAzureCli => Some(self.use_azure_cli.to_string()),
382            AzureConfigKey::SkipSignature => Some(self.skip_signature.to_string()),
383            AzureConfigKey::Client(key) => self.client_options.get_config_value(key),
384        }
385    }
386
387    /// Set the Azure Account
388    pub fn with_account(mut self, account: impl Into<String>) -> Self {
389        self.account_name = Some(account.into());
390        self
391    }
392
393    /// Set the Azure Access Key
394    pub fn with_access_key(mut self, access_key: impl Into<String>) -> Self {
395        self.access_key = Some(access_key.into());
396        self
397    }
398
399    /// Set a static bearer token to be used for authorizing requests
400    pub fn with_bearer_token_authorization(mut self, bearer_token: impl Into<String>) -> Self {
401        self.bearer_token = Some(bearer_token.into());
402        self
403    }
404
405    /// Set a client secret used for client secret authorization
406    pub fn with_client_secret_authorization(
407        mut self,
408        client_id: impl Into<String>,
409        client_secret: impl Into<String>,
410        tenant_id: impl Into<String>,
411    ) -> Self {
412        self.client_id = Some(client_id.into());
413        self.client_secret = Some(client_secret.into());
414        self.tenant_id = Some(tenant_id.into());
415        self
416    }
417
418    /// Sets the client id for use in client secret or k8s federated credential flow
419    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
420        self.client_id = Some(client_id.into());
421        self
422    }
423
424    /// Sets the client secret for use in client secret flow
425    pub fn with_client_secret(mut self, client_secret: impl Into<String>) -> Self {
426        self.client_secret = Some(client_secret.into());
427        self
428    }
429
430    /// Sets the tenant id for use in client secret or k8s federated credential flow
431    pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
432        self.tenant_id = Some(tenant_id.into());
433        self
434    }
435
436    /// Set the OAuth scope for client credentials flow
437    pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
438        self.scope = Some(scope.into());
439        self
440    }
441
442    /// Set the credential provider overriding any other options
443    pub fn with_credentials(mut self, credentials: AzureCredentialProvider) -> Self {
444        self.credentials = Some(credentials);
445        self
446    }
447
448    /// Override the endpoint used to communicate with Azure
449    pub fn with_endpoint(mut self, endpoint: String) -> Self {
450        self.endpoint = Some(endpoint);
451        self
452    }
453
454    /// Sets what protocol is allowed
455    pub fn with_allow_http(mut self, allow_http: bool) -> Self {
456        self.client_options = self.client_options.with_allow_http(allow_http);
457        self
458    }
459
460    /// Sets an alternative authority host for OAuth based authorization
461    ///
462    /// Defaults to <https://login.microsoftonline.com>
463    pub fn with_authority_host(mut self, authority_host: impl Into<String>) -> Self {
464        self.authority_host = Some(authority_host.into());
465        self
466    }
467
468    /// Set the retry configuration
469    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
470        self.retry_config = retry_config;
471        self
472    }
473
474    /// Set the proxy_url to be used by the underlying client
475    pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
476        self.client_options = self.client_options.with_proxy_url(proxy_url);
477        self
478    }
479
480    /// Set a trusted proxy CA certificate
481    pub fn with_proxy_ca_certificate(mut self, proxy_ca_certificate: impl Into<String>) -> Self {
482        self.client_options = self
483            .client_options
484            .with_proxy_ca_certificate(proxy_ca_certificate);
485        self
486    }
487
488    /// Set a list of hosts to exclude from proxy connections
489    pub fn with_proxy_excludes(mut self, proxy_excludes: impl Into<String>) -> Self {
490        self.client_options = self.client_options.with_proxy_excludes(proxy_excludes);
491        self
492    }
493
494    /// Sets the client options, overriding any already set
495    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
496        self.client_options = options;
497        self
498    }
499
500    /// Sets the endpoint for acquiring managed identity token
501    pub fn with_msi_endpoint(mut self, msi_endpoint: impl Into<String>) -> Self {
502        self.msi_endpoint = Some(msi_endpoint.into());
503        self
504    }
505
506    /// Sets a file path for acquiring azure federated identity token in k8s
507    pub fn with_federated_token_file(mut self, federated_token_file: impl Into<String>) -> Self {
508        self.federated_token_file = Some(federated_token_file.into());
509        self
510    }
511
512    /// Set if the Azure Cli should be used for acquiring access token
513    pub fn with_use_azure_cli(mut self, use_azure_cli: bool) -> Self {
514        self.use_azure_cli = use_azure_cli.into();
515        self
516    }
517
518    /// If enabled, requests will not be signed
519    pub fn with_skip_signature(mut self, skip_signature: bool) -> Self {
520        self.skip_signature = skip_signature.into();
521        self
522    }
523
524    /// Build an [`AzureConfig`] from the provided values, consuming `self`.
525    /// Build an [`AzureConfig`] from the provided values, consuming `self`.
526    ///
527    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
528    /// will be spawned on the given runtime handle.
529    pub fn build(self, runtime: Option<&Handle>) -> Result<AzureConfig> {
530        let static_creds = |credential: AzureCredential| -> AzureCredentialProvider {
531            Arc::new(StaticCredentialProvider::new(credential))
532        };
533
534        let auth = if let Some(credential) = self.credentials {
535            credential
536        } else if let Some(bearer_token) = self.bearer_token {
537            static_creds(AzureCredential::BearerToken(bearer_token))
538        } else if let (Some(client_id), Some(tenant_id), Some(federated_token_file)) =
539            (&self.client_id, &self.tenant_id, self.federated_token_file)
540        {
541            let client_credential = WorkloadIdentityOAuthProvider::new(
542                client_id,
543                federated_token_file,
544                tenant_id,
545                self.authority_host,
546            );
547            let client = self.client_options.client()?;
548            let service = make_service(client.clone(), runtime);
549            Arc::new(TokenCredentialProvider::new(
550                client_credential,
551                client,
552                service,
553                self.retry_config.clone(),
554            )) as _
555        } else if let (Some(client_id), Some(client_secret), Some(tenant_id)) =
556            (&self.client_id, self.client_secret, &self.tenant_id)
557        {
558            let scope = self.scope.ok_or(Error::MissingScope)?;
559            let client_credential = ClientSecretOAuthProvider::new_with_scope(
560                client_id.clone(),
561                client_secret,
562                tenant_id,
563                self.authority_host,
564                &scope,
565            );
566            let client = self.client_options.client()?;
567            let service = make_service(client.clone(), runtime);
568            Arc::new(TokenCredentialProvider::new(
569                client_credential,
570                client,
571                service,
572                self.retry_config.clone(),
573            )) as _
574        } else if self.use_azure_cli.get()? {
575            Arc::new(AzureCliCredential::new()) as _
576        } else {
577            let msi_credential = ImdsManagedIdentityProvider::new(
578                self.client_id,
579                self.object_id,
580                self.msi_resource_id,
581                self.msi_endpoint,
582            );
583            let client = self.client_options.metadata_client()?;
584            let service = make_service(client.clone(), runtime);
585            Arc::new(TokenCredentialProvider::new(
586                msi_credential,
587                client,
588                service,
589                self.retry_config.clone(),
590            )) as _
591        };
592
593        Ok(AzureConfig {
594            skip_signature: self.skip_signature.get()?,
595            retry_config: self.retry_config,
596            client_options: self.client_options,
597            credentials: auth,
598        })
599    }
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605    use std::collections::HashMap;
606
607    #[test]
608    fn azure_test_config_from_map() {
609        let azure_client_id = "object_store:fake_access_key_id";
610        let azure_storage_account_name = "object_store:fake_secret_key";
611        let azure_storage_token = "object_store:fake_default_region";
612        let options = HashMap::from([
613            ("azure_client_id", azure_client_id),
614            ("azure_storage_account_name", azure_storage_account_name),
615            ("azure_storage_token", azure_storage_token),
616        ]);
617
618        let builder = options
619            .into_iter()
620            .fold(AzureBuilder::new(), |builder, (key, value)| {
621                builder.with_config(key.parse().unwrap(), value)
622            });
623        assert_eq!(builder.client_id.unwrap(), azure_client_id);
624        assert_eq!(builder.account_name.unwrap(), azure_storage_account_name);
625        assert_eq!(builder.bearer_token.unwrap(), azure_storage_token);
626    }
627
628    #[test]
629    fn azure_test_client_opts() {
630        let key = "AZURE_PROXY_URL";
631        if let Ok(config_key) = key.to_ascii_lowercase().parse() {
632            assert_eq!(
633                AzureConfigKey::Client(ClientConfigKey::ProxyUrl),
634                config_key
635            );
636        } else {
637            panic!("{key} not propagated as ClientConfigKey");
638        }
639    }
640}
olai_http/azure/builder.rs

olai_http/azure/
builder.rs
