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                && key.starts_with("AZURE_")
326                && let Ok(config_key) = key.to_ascii_lowercase().parse()
327            {
328                builder = builder.with_config(config_key, value);
329            }
330        }
331
332        if let Ok(text) = std::env::var(MSI_ENDPOINT_ENV_KEY) {
333            builder = builder.with_msi_endpoint(text);
334        }
335
336        builder
337    }
338
339    /// Set an option on the builder via a key - value pair.
340    pub fn with_config(mut self, key: AzureConfigKey, value: impl Into<String>) -> Self {
341        match key {
342            AzureConfigKey::AccessKey => self.access_key = Some(value.into()),
343            AzureConfigKey::AccountName => self.account_name = Some(value.into()),
344            AzureConfigKey::ClientId => self.client_id = Some(value.into()),
345            AzureConfigKey::ClientSecret => self.client_secret = Some(value.into()),
346            AzureConfigKey::AuthorityId => self.tenant_id = Some(value.into()),
347            AzureConfigKey::AuthorityHost => self.authority_host = Some(value.into()),
348            AzureConfigKey::Scope => self.scope = Some(value.into()),
349            AzureConfigKey::Token => self.bearer_token = Some(value.into()),
350            AzureConfigKey::MsiEndpoint => self.msi_endpoint = Some(value.into()),
351            AzureConfigKey::ObjectId => self.object_id = Some(value.into()),
352            AzureConfigKey::MsiResourceId => self.msi_resource_id = Some(value.into()),
353            AzureConfigKey::FederatedTokenFile => self.federated_token_file = Some(value.into()),
354            AzureConfigKey::UseAzureCli => self.use_azure_cli.parse(value),
355            AzureConfigKey::SkipSignature => self.skip_signature.parse(value),
356            AzureConfigKey::Endpoint => self.endpoint = Some(value.into()),
357            AzureConfigKey::Client(key) => {
358                self.client_options = self.client_options.with_config(key, value)
359            }
360        };
361        self
362    }
363
364    /// Get config value via a [`AzureConfigKey`].
365    pub fn get_config_value(&self, key: &AzureConfigKey) -> Option<String> {
366        match key {
367            AzureConfigKey::AccountName => self.account_name.clone(),
368            AzureConfigKey::AccessKey => self.access_key.clone(),
369            AzureConfigKey::ClientId => self.client_id.clone(),
370            AzureConfigKey::ClientSecret => self.client_secret.clone(),
371            AzureConfigKey::AuthorityId => self.tenant_id.clone(),
372            AzureConfigKey::AuthorityHost => self.authority_host.clone(),
373            AzureConfigKey::Scope => self.scope.clone(),
374            AzureConfigKey::Token => self.bearer_token.clone(),
375            AzureConfigKey::Endpoint => self.endpoint.clone(),
376            AzureConfigKey::MsiEndpoint => self.msi_endpoint.clone(),
377            AzureConfigKey::ObjectId => self.object_id.clone(),
378            AzureConfigKey::MsiResourceId => self.msi_resource_id.clone(),
379            AzureConfigKey::FederatedTokenFile => self.federated_token_file.clone(),
380            AzureConfigKey::UseAzureCli => Some(self.use_azure_cli.to_string()),
381            AzureConfigKey::SkipSignature => Some(self.skip_signature.to_string()),
382            AzureConfigKey::Client(key) => self.client_options.get_config_value(key),
383        }
384    }
385
386    /// Set the Azure Account
387    pub fn with_account(mut self, account: impl Into<String>) -> Self {
388        self.account_name = Some(account.into());
389        self
390    }
391
392    /// Set the Azure Access Key
393    pub fn with_access_key(mut self, access_key: impl Into<String>) -> Self {
394        self.access_key = Some(access_key.into());
395        self
396    }
397
398    /// Set a static bearer token to be used for authorizing requests
399    pub fn with_bearer_token_authorization(mut self, bearer_token: impl Into<String>) -> Self {
400        self.bearer_token = Some(bearer_token.into());
401        self
402    }
403
404    /// Set a client secret used for client secret authorization
405    pub fn with_client_secret_authorization(
406        mut self,
407        client_id: impl Into<String>,
408        client_secret: impl Into<String>,
409        tenant_id: impl Into<String>,
410    ) -> Self {
411        self.client_id = Some(client_id.into());
412        self.client_secret = Some(client_secret.into());
413        self.tenant_id = Some(tenant_id.into());
414        self
415    }
416
417    /// Sets the client id for use in client secret or k8s federated credential flow
418    pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
419        self.client_id = Some(client_id.into());
420        self
421    }
422
423    /// Sets the client secret for use in client secret flow
424    pub fn with_client_secret(mut self, client_secret: impl Into<String>) -> Self {
425        self.client_secret = Some(client_secret.into());
426        self
427    }
428
429    /// Sets the tenant id for use in client secret or k8s federated credential flow
430    pub fn with_tenant_id(mut self, tenant_id: impl Into<String>) -> Self {
431        self.tenant_id = Some(tenant_id.into());
432        self
433    }
434
435    /// Set the OAuth scope for client credentials flow
436    pub fn with_scope(mut self, scope: impl Into<String>) -> Self {
437        self.scope = Some(scope.into());
438        self
439    }
440
441    /// Set the credential provider overriding any other options
442    pub fn with_credentials(mut self, credentials: AzureCredentialProvider) -> Self {
443        self.credentials = Some(credentials);
444        self
445    }
446
447    /// Override the endpoint used to communicate with Azure
448    pub fn with_endpoint(mut self, endpoint: String) -> Self {
449        self.endpoint = Some(endpoint);
450        self
451    }
452
453    /// Sets what protocol is allowed
454    pub fn with_allow_http(mut self, allow_http: bool) -> Self {
455        self.client_options = self.client_options.with_allow_http(allow_http);
456        self
457    }
458
459    /// Sets an alternative authority host for OAuth based authorization
460    ///
461    /// Defaults to <https://login.microsoftonline.com>
462    pub fn with_authority_host(mut self, authority_host: impl Into<String>) -> Self {
463        self.authority_host = Some(authority_host.into());
464        self
465    }
466
467    /// Set the retry configuration
468    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
469        self.retry_config = retry_config;
470        self
471    }
472
473    /// Set the proxy_url to be used by the underlying client
474    pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
475        self.client_options = self.client_options.with_proxy_url(proxy_url);
476        self
477    }
478
479    /// Set a trusted proxy CA certificate
480    pub fn with_proxy_ca_certificate(mut self, proxy_ca_certificate: impl Into<String>) -> Self {
481        self.client_options = self
482            .client_options
483            .with_proxy_ca_certificate(proxy_ca_certificate);
484        self
485    }
486
487    /// Set a list of hosts to exclude from proxy connections
488    pub fn with_proxy_excludes(mut self, proxy_excludes: impl Into<String>) -> Self {
489        self.client_options = self.client_options.with_proxy_excludes(proxy_excludes);
490        self
491    }
492
493    /// Sets the client options, overriding any already set
494    pub fn with_client_options(mut self, options: ClientOptions) -> Self {
495        self.client_options = options;
496        self
497    }
498
499    /// Sets the endpoint for acquiring managed identity token
500    pub fn with_msi_endpoint(mut self, msi_endpoint: impl Into<String>) -> Self {
501        self.msi_endpoint = Some(msi_endpoint.into());
502        self
503    }
504
505    /// Sets a file path for acquiring azure federated identity token in k8s
506    pub fn with_federated_token_file(mut self, federated_token_file: impl Into<String>) -> Self {
507        self.federated_token_file = Some(federated_token_file.into());
508        self
509    }
510
511    /// Set if the Azure Cli should be used for acquiring access token
512    pub fn with_use_azure_cli(mut self, use_azure_cli: bool) -> Self {
513        self.use_azure_cli = use_azure_cli.into();
514        self
515    }
516
517    /// If enabled, requests will not be signed
518    pub fn with_skip_signature(mut self, skip_signature: bool) -> Self {
519        self.skip_signature = skip_signature.into();
520        self
521    }
522
523    /// Build an [`AzureConfig`] from the provided values, consuming `self`.
524    /// Build an [`AzureConfig`] from the provided values, consuming `self`.
525    ///
526    /// If `runtime` is provided, all HTTP I/O (including credential refresh)
527    /// will be spawned on the given runtime handle.
528    pub fn build(self, runtime: Option<&Handle>) -> Result<AzureConfig> {
529        let static_creds = |credential: AzureCredential| -> AzureCredentialProvider {
530            Arc::new(StaticCredentialProvider::new(credential))
531        };
532
533        let auth = if let Some(credential) = self.credentials {
534            credential
535        } else if let Some(bearer_token) = self.bearer_token {
536            static_creds(AzureCredential::BearerToken(bearer_token))
537        } else if let (Some(client_id), Some(tenant_id), Some(federated_token_file)) =
538            (&self.client_id, &self.tenant_id, self.federated_token_file)
539        {
540            let client_credential = WorkloadIdentityOAuthProvider::new(
541                client_id,
542                federated_token_file,
543                tenant_id,
544                self.authority_host,
545            );
546            let client = self.client_options.client()?;
547            let service = make_service(client.clone(), runtime);
548            Arc::new(TokenCredentialProvider::new(
549                client_credential,
550                client,
551                service,
552                self.retry_config.clone(),
553            )) as _
554        } else if let (Some(client_id), Some(client_secret), Some(tenant_id)) =
555            (&self.client_id, self.client_secret, &self.tenant_id)
556        {
557            let scope = self.scope.ok_or(Error::MissingScope)?;
558            let client_credential = ClientSecretOAuthProvider::new_with_scope(
559                client_id.clone(),
560                client_secret,
561                tenant_id,
562                self.authority_host,
563                &scope,
564            );
565            let client = self.client_options.client()?;
566            let service = make_service(client.clone(), runtime);
567            Arc::new(TokenCredentialProvider::new(
568                client_credential,
569                client,
570                service,
571                self.retry_config.clone(),
572            )) as _
573        } else if self.use_azure_cli.get()? {
574            Arc::new(AzureCliCredential::new()) as _
575        } else {
576            let msi_credential = ImdsManagedIdentityProvider::new(
577                self.client_id,
578                self.object_id,
579                self.msi_resource_id,
580                self.msi_endpoint,
581            );
582            let client = self.client_options.metadata_client()?;
583            let service = make_service(client.clone(), runtime);
584            Arc::new(TokenCredentialProvider::new(
585                msi_credential,
586                client,
587                service,
588                self.retry_config.clone(),
589            )) as _
590        };
591
592        Ok(AzureConfig {
593            skip_signature: self.skip_signature.get()?,
594            retry_config: self.retry_config,
595            client_options: self.client_options,
596            credentials: auth,
597        })
598    }
599}
600
601#[cfg(test)]
602mod tests {
603    use super::*;
604    use std::collections::HashMap;
605
606    #[test]
607    fn azure_test_config_from_map() {
608        let azure_client_id = "object_store:fake_access_key_id";
609        let azure_storage_account_name = "object_store:fake_secret_key";
610        let azure_storage_token = "object_store:fake_default_region";
611        let options = HashMap::from([
612            ("azure_client_id", azure_client_id),
613            ("azure_storage_account_name", azure_storage_account_name),
614            ("azure_storage_token", azure_storage_token),
615        ]);
616
617        let builder = options
618            .into_iter()
619            .fold(AzureBuilder::new(), |builder, (key, value)| {
620                builder.with_config(key.parse().unwrap(), value)
621            });
622        assert_eq!(builder.client_id.unwrap(), azure_client_id);
623        assert_eq!(builder.account_name.unwrap(), azure_storage_account_name);
624        assert_eq!(builder.bearer_token.unwrap(), azure_storage_token);
625    }
626
627    #[test]
628    fn azure_test_client_opts() {
629        let key = "AZURE_PROXY_URL";
630        if let Ok(config_key) = key.to_ascii_lowercase().parse() {
631            assert_eq!(
632                AzureConfigKey::Client(ClientConfigKey::ProxyUrl),
633                config_key
634            );
635        } else {
636            panic!("{key} not propagated as ClientConfigKey");
637        }
638    }
639}
olai_http/azure/builder.rs

olai_http/azure/
builder.rs
