azure_identity/token_credentials/
default_credentials.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
#[cfg(not(target_arch = "wasm32"))]
use crate::AzureCliCredential;
use crate::{
    timeout::TimeoutExt, token_credentials::cache::TokenCache, AppServiceManagedIdentityCredential,
    EnvironmentCredential, TokenCredentialOptions, VirtualMachineManagedIdentityCredential,
};
use azure_core::{
    auth::{AccessToken, TokenCredential},
    error::{Error, ErrorKind, ResultExt},
};
use std::{sync::Arc, time::Duration};

/// Provides a mechanism of selectively disabling credentials used for a `DefaultAzureCredential` instance
pub struct DefaultAzureCredentialBuilder {
    options: TokenCredentialOptions,
    include_environment_credential: bool,
    include_app_service_managed_identity_credential: bool,
    include_virtual_machine_managed_identity_credential: bool,
    #[cfg(not(target_arch = "wasm32"))]
    include_azure_cli_credential: bool,
}

impl Default for DefaultAzureCredentialBuilder {
    fn default() -> Self {
        Self {
            options: TokenCredentialOptions::default(),
            include_environment_credential: true,
            include_app_service_managed_identity_credential: true,
            include_virtual_machine_managed_identity_credential: true,
            #[cfg(not(target_arch = "wasm32"))]
            include_azure_cli_credential: true,
        }
    }
}

impl DefaultAzureCredentialBuilder {
    /// Create a new `DefaultAzureCredentialBuilder`
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_options(&mut self, options: impl Into<TokenCredentialOptions>) -> &mut Self {
        self.options = options.into();
        self
    }

    /// Exclude using any environment credential
    pub fn exclude_environment_credential(&mut self) -> &mut Self {
        self.include_environment_credential = false;
        self
    }

    /// Exclude using any managed identity credential
    pub fn exclude_managed_identity_credential(&mut self) -> &mut Self {
        self.include_app_service_managed_identity_credential = false;
        self.include_virtual_machine_managed_identity_credential = false;
        self
    }

    /// Exclude using virtual machine managed identity credential
    pub fn exclude_virtual_machine_managed_identity_credential(&mut self) -> &mut Self {
        self.include_virtual_machine_managed_identity_credential = false;
        self
    }

    /// Include using virtual machine managed identity credential
    pub fn include_virtual_machine_managed_identity_credential(&mut self) -> &mut Self {
        self.include_virtual_machine_managed_identity_credential = true;
        self
    }

    /// Include using app service managed identity credential
    pub fn include_app_service_managed_identity_credential(&mut self) -> &mut Self {
        self.include_app_service_managed_identity_credential = true;
        self
    }

    /// Exclude using credential from the cli
    #[cfg(not(target_arch = "wasm32"))]
    pub fn exclude_azure_cli_credential(&mut self) -> &mut Self {
        self.include_azure_cli_credential = false;
        self
    }

    /// Get a list of the credential types to include.
    fn included(&self) -> Vec<DefaultAzureCredentialType> {
        let mut sources = Vec::new();
        if self.include_environment_credential {
            sources.push(DefaultAzureCredentialType::Environment);
        }
        if self.include_app_service_managed_identity_credential {
            sources.push(DefaultAzureCredentialType::AppService);
        }
        if self.include_virtual_machine_managed_identity_credential {
            sources.push(DefaultAzureCredentialType::VirtualMachine);
        }
        #[cfg(not(target_arch = "wasm32"))]
        if self.include_azure_cli_credential {
            sources.push(DefaultAzureCredentialType::AzureCli);
        }
        sources
    }

    /// Creates a list of `TokenCredential` instances from the included credential types.
    /// The credentials created successfully are used as sources for getting a token.
    fn create_sources(
        &self,
        included: &Vec<DefaultAzureCredentialType>,
    ) -> azure_core::Result<Vec<DefaultAzureCredentialKind>> {
        let mut sources = Vec::<DefaultAzureCredentialKind>::with_capacity(included.len());
        let mut errors = Vec::new();
        for source in included {
            match source {
                DefaultAzureCredentialType::Environment => {
                    match EnvironmentCredential::create(self.options.clone()) {
                        Ok(credential) => {
                            sources.push(DefaultAzureCredentialKind::Environment(credential))
                        }
                        Err(error) => errors.push(error),
                    }
                }
                DefaultAzureCredentialType::AppService => {
                    match AppServiceManagedIdentityCredential::create(self.options.clone()) {
                        Ok(credential) => {
                            sources.push(DefaultAzureCredentialKind::AppService(credential))
                        }
                        Err(error) => errors.push(error),
                    }
                }
                DefaultAzureCredentialType::VirtualMachine => {
                    sources.push(DefaultAzureCredentialKind::VirtualMachine(
                        VirtualMachineManagedIdentityCredential::new(self.options.clone()),
                    ));
                }
                #[cfg(not(target_arch = "wasm32"))]
                DefaultAzureCredentialType::AzureCli => {
                    if let Ok(credential) = AzureCliCredential::create() {
                        sources.push(DefaultAzureCredentialKind::AzureCli(credential));
                    }
                }
            }
        }
        if sources.is_empty() {
            return Err(Error::with_message(ErrorKind::Credential, || {
                format!(
                    "No credential sources were available to be used for authentication.\n{}",
                    format_aggregate_error(&errors)
                )
            }));
        }
        Ok(sources)
    }

    /// Create a `DefaultAzureCredential` from this builder.
    pub fn build(&self) -> azure_core::Result<DefaultAzureCredential> {
        let included = self.included();
        let sources = self.create_sources(&included)?;
        Ok(DefaultAzureCredential::with_sources(sources))
    }
}

/// Types that may be enabled for use by `DefaultAzureCredential`.
#[derive(Debug, PartialEq)]
enum DefaultAzureCredentialType {
    Environment,
    AppService,
    VirtualMachine,
    #[cfg(not(target_arch = "wasm32"))]
    AzureCli,
}

/// Types of `TokenCredential` supported by `DefaultAzureCredential`
#[derive(Debug)]
pub(crate) enum DefaultAzureCredentialKind {
    /// `TokenCredential` from environment variable.
    Environment(EnvironmentCredential),
    /// `TokenCredential` from managed identity that has been assigned to an App Service.
    AppService(AppServiceManagedIdentityCredential),
    /// `TokenCredential` from managed identity that has been assigned to a virtual machine.
    VirtualMachine(VirtualMachineManagedIdentityCredential),
    #[cfg(not(target_arch = "wasm32"))]
    /// `TokenCredential` from Azure CLI.
    AzureCli(AzureCliCredential),
}

#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
impl TokenCredential for DefaultAzureCredentialKind {
    async fn get_token(&self, scopes: &[&str]) -> azure_core::Result<AccessToken> {
        match self {
            DefaultAzureCredentialKind::Environment(credential) => {
                credential.get_token(scopes).await.context(
                    ErrorKind::Credential,
                    "error getting environment credential",
                )
            }
            DefaultAzureCredentialKind::AppService(credential) => {
                credential.get_token(scopes).await.context(
                    ErrorKind::Credential,
                    "error getting managed identity credential for App Service",
                )
            }
            DefaultAzureCredentialKind::VirtualMachine(credential) => {
                // IMSD timeout is only limited to 1 second when used in DefaultAzureCredential
                credential
                    .get_token(scopes)
                    .timeout(Duration::from_secs(1))
                    .await
                    .context(
                        ErrorKind::Credential,
                        "getting virtual machine managed identity credential timed out",
                    )?
                    .context(
                        ErrorKind::Credential,
                        "error getting virtual machine managed identity credential",
                    )
            }
            #[cfg(not(target_arch = "wasm32"))]
            DefaultAzureCredentialKind::AzureCli(credential) => {
                credential.get_token(scopes).await.context(
                    ErrorKind::Credential,
                    "error getting token credential from Azure CLI",
                )
            }
        }
    }

    /// Clear the credential's cache.
    async fn clear_cache(&self) -> azure_core::Result<()> {
        match self {
            DefaultAzureCredentialKind::Environment(credential) => credential.clear_cache().await,
            DefaultAzureCredentialKind::AppService(credential) => credential.clear_cache().await,
            DefaultAzureCredentialKind::VirtualMachine(credential) => {
                credential.clear_cache().await
            }
            #[cfg(not(target_arch = "wasm32"))]
            DefaultAzureCredentialKind::AzureCli(credential) => credential.clear_cache().await,
        }
    }
}

/// Provides a default `TokenCredential` authentication flow for applications that will be deployed to Azure.
///
/// The following credential types if enabled will be tried, in order:
/// - `EnvironmentCredential`
/// - `ManagedIdentityCredential`
/// - `AzureCliCredential`
///
/// Consult the documentation of these credential types for more information on how they attempt authentication.
#[derive(Debug)]
pub struct DefaultAzureCredential {
    sources: Vec<DefaultAzureCredentialKind>,
    cache: TokenCache,
}

impl DefaultAzureCredential {
    pub fn create(options: TokenCredentialOptions) -> azure_core::Result<DefaultAzureCredential> {
        DefaultAzureCredentialBuilder::default()
            .with_options(options)
            .build()
    }

    /// Creates a `DefaultAzureCredential` with specified sources.
    fn with_sources(sources: Vec<DefaultAzureCredentialKind>) -> Self {
        DefaultAzureCredential {
            sources,
            cache: TokenCache::new(),
        }
    }

    /// Try to fetch a token using each of the credential sources until one succeeds
    async fn get_token(&self, scopes: &[&str]) -> azure_core::Result<AccessToken> {
        let mut errors = Vec::new();
        for source in &self.sources {
            let token_res = source.get_token(scopes).await;

            match token_res {
                Ok(token) => return Ok(token),
                Err(error) => errors.push(error),
            }
        }
        Err(Error::with_message(ErrorKind::Credential, || {
            format!(
                "Multiple errors were encountered while attempting to authenticate:\n{}",
                format_aggregate_error(&errors)
            )
        }))
    }
}

/// Creates a new `DefaultAzureCredential` with the default options.
pub fn create_default_credential() -> azure_core::Result<Arc<dyn TokenCredential>> {
    DefaultAzureCredentialBuilder::default()
        .build()
        .map(|cred| Arc::new(cred) as Arc<dyn TokenCredential>)
}

#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
impl TokenCredential for DefaultAzureCredential {
    async fn get_token(&self, scopes: &[&str]) -> azure_core::Result<AccessToken> {
        self.cache.get_token(scopes, self.get_token(scopes)).await
    }

    /// Clear the credential's cache.
    async fn clear_cache(&self) -> azure_core::Result<()> {
        // clear the internal cache as well as each of the underlying providers
        self.cache.clear().await?;

        for source in &self.sources {
            source.clear_cache().await?;
        }

        Ok(())
    }
}

fn format_aggregate_error(errors: &[Error]) -> String {
    use std::error::Error;
    errors
        .iter()
        .map(|e| {
            let mut current: Option<&dyn Error> = Some(e);
            let mut stack = vec![];
            while let Some(err) = current.take() {
                stack.push(err.to_string());
                current = err.source();
            }
            stack.join(" - ")
        })
        .collect::<Vec<String>>()
        .join("\n")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builder_included_credential_flags() {
        let builder = DefaultAzureCredentialBuilder::new();
        #[cfg(not(target_arch = "wasm32"))]
        assert!(builder.include_azure_cli_credential);
        assert!(builder.include_environment_credential);
        assert!(builder.include_app_service_managed_identity_credential);
        assert!(builder.include_virtual_machine_managed_identity_credential);

        #[cfg(not(target_arch = "wasm32"))]
        {
            let mut builder = DefaultAzureCredentialBuilder::new();
            builder.exclude_azure_cli_credential();
            assert!(!builder.include_azure_cli_credential);
            assert!(builder.include_environment_credential);
            assert!(builder.include_app_service_managed_identity_credential);
            assert!(builder.include_virtual_machine_managed_identity_credential);
        }

        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_environment_credential();
        #[cfg(not(target_arch = "wasm32"))]
        assert!(builder.include_azure_cli_credential);
        assert!(!builder.include_environment_credential);
        assert!(builder.include_app_service_managed_identity_credential);
        assert!(builder.include_virtual_machine_managed_identity_credential);

        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_managed_identity_credential();
        #[cfg(not(target_arch = "wasm32"))]
        assert!(builder.include_azure_cli_credential);
        assert!(builder.include_environment_credential);
        assert!(!builder.include_app_service_managed_identity_credential);
        assert!(!builder.include_virtual_machine_managed_identity_credential);
    }

    #[test]
    /// test default included credential types
    fn test_default_included_credential_types() {
        let builder = DefaultAzureCredentialBuilder::new();
        assert_eq!(
            builder.included(),
            vec![
                DefaultAzureCredentialType::Environment,
                DefaultAzureCredentialType::AppService,
                DefaultAzureCredentialType::VirtualMachine,
                DefaultAzureCredentialType::AzureCli,
            ]
        );
    }

    /// test excluding virtual machine managed identity credential
    #[test]
    fn test_exclude_virtual_machine_managed_identity_credential() {
        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_virtual_machine_managed_identity_credential();
        assert_eq!(
            builder.included(),
            vec![
                DefaultAzureCredentialType::Environment,
                DefaultAzureCredentialType::AppService,
                DefaultAzureCredentialType::AzureCli,
            ]
        );
    }

    /// test excluding environment credential
    #[test]
    fn test_exclude_environment_credential() -> azure_core::Result<()> {
        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_environment_credential();
        assert_eq!(
            builder.included(),
            vec![
                DefaultAzureCredentialType::AppService,
                DefaultAzureCredentialType::VirtualMachine,
                DefaultAzureCredentialType::AzureCli,
            ]
        );
        Ok(())
    }

    /// test excluding azure cli credential
    #[test]
    fn test_exclude_azure_cli_credential() {
        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_azure_cli_credential();
        assert_eq!(
            builder.included(),
            vec![
                DefaultAzureCredentialType::Environment,
                DefaultAzureCredentialType::AppService,
                DefaultAzureCredentialType::VirtualMachine,
            ]
        );
    }

    /// test exluding managed identity credential
    #[test]
    fn test_exclude_managed_identity_credential() {
        let mut builder = DefaultAzureCredentialBuilder::new();
        builder.exclude_managed_identity_credential();
        assert_eq!(
            builder.included(),
            vec![
                DefaultAzureCredentialType::Environment,
                DefaultAzureCredentialType::AzureCli,
            ]
        );
    }
}