Live ACI and IMDS managed identity tests (#2775)

chlowell · web-flow · commit 5127760cec18 · 2025-07-10T09:33:52.000-07:00
diff --git a/.vscode/cspell.json b/.vscode/cspell.json
@@ -5,6 +5,7 @@
   "ignorePaths": [
     "**/test-resources.bicep",
     "**/test-resources.json",
+    "**/test-resources-post.ps1",
     "**/assets.json",
     ".config",
     ".devcontainer/devcontainer.json",
diff --git a/sdk/identity/.dict.txt b/sdk/identity/.dict.txt
@@ -6,6 +6,8 @@ clientcertificate
 clientsecret
 cloudshell
 imds
+LINUXPOOL
+LINUXVMIMAGE
 managedidentity
 msal
 replacen
diff --git a/sdk/identity/azure_identity/managed-identity-matrix.json b/sdk/identity/azure_identity/managed-identity-matrix.json
@@ -0,0 +1,17 @@
+{
+    "include": [
+        {
+            "Agent": {
+                "msi_image": {
+                    "ArmTemplateParameters": "@{deployResources = $true}",
+                    "OSVmImage": "env:LINUXVMIMAGE",
+                    "Pool": "env:LINUXPOOL"
+                }
+            },
+            "IDENTITY_IMDS_AVAILABLE": "1",
+            "RustToolchainName": [
+                "stable"
+            ]
+        }
+    ]
+}
diff --git a/sdk/identity/azure_identity/src/managed_identity_credential.rs b/sdk/identity/azure_identity/src/managed_identity_credential.rs
@@ -153,15 +153,42 @@ mod tests {
     use crate::env::Env;
     use crate::tests::{LIVE_TEST_RESOURCE, LIVE_TEST_SCOPES};
     use azure_core::http::headers::Headers;
-    use azure_core::http::{Method, RawResponse, Request, StatusCode};
+    use azure_core::http::{Method, RawResponse, Request, StatusCode, Url};
+    use azure_core::time::OffsetDateTime;
     use azure_core::Bytes;
-    use azure_core_test::http::MockHttpClient;
+    use azure_core_test::{http::MockHttpClient, recorded};
     use futures::FutureExt;
+    use std::env;
     use std::sync::atomic::{AtomicUsize, Ordering};
     use std::time::{SystemTime, UNIX_EPOCH};
 
     const EXPIRES_ON: &str = "EXPIRES_ON";
 
+    async fn run_deployed_test(
+        authority: &str,
+        storage_name: &str,
+        id: Option<UserAssignedId>,
+    ) -> azure_core::Result<()> {
+        let id_param = id.map_or("".to_string(), |id| match id {
+            UserAssignedId::ClientId(id) => format!("client-id={id}&"),
+            UserAssignedId::ObjectId(id) => format!("object-id={id}&"),
+            UserAssignedId::ResourceId(id) => format!("resource-id={id}&"),
+        });
+        let url = format!(
+            "http://{authority}/api?test=managed-identity&{id_param}storage-name={storage_name}"
+        );
+        let u = Url::parse(&url).expect("invalid URL");
+        let client = azure_core::http::new_http_client();
+        let req = Request::new(u, Method::Get);
+
+        let res = client.execute_request(&req).await.expect("request failed");
+        let status = res.status();
+        let body = res.into_body().collect_string().await?;
+        assert_eq!(StatusCode::Ok, status, "Test app responded with '{body}'");
+
+        Ok(())
+    }
+
     async fn run_supported_source_test(
         env: Env,
         options: Option<ManagedIdentityCredentialOptions>,
@@ -256,6 +283,27 @@ mod tests {
         );
     }
 
+    #[recorded::test(live)]
+    async fn aci_user_assigned_live() -> azure_core::Result<()> {
+        if env::var("CI_HAS_DEPLOYED_RESOURCES").is_err() {
+            println!("Skipped: ACI live tests require deployed resources");
+            return Ok(());
+        }
+        let ip = env::var("IDENTITY_ACI_IP_USER_ASSIGNED").expect("IDENTITY_ACI_IP_USER_ASSIGNED");
+        let storage_name = env::var("IDENTITY_STORAGE_NAME_USER_ASSIGNED")
+            .expect("IDENTITY_STORAGE_NAME_USER_ASSIGNED");
+        let client_id = env::var("IDENTITY_USER_ASSIGNED_IDENTITY_CLIENT_ID")
+            .expect("IDENTITY_USER_ASSIGNED_IDENTITY_CLIENT_ID");
+        run_deployed_test(
+            &format!("{}:8080", ip),
+            &storage_name,
+            Some(UserAssignedId::ClientId(client_id)),
+        )
+        .await?;
+
+        Ok(())
+    }
+
     async fn run_app_service_test(options: Option<ManagedIdentityCredentialOptions>) {
         let endpoint = "http://localhost/metadata/identity/oauth2/token";
         let x_id_header = "x-id-header";
@@ -369,6 +417,27 @@ mod tests {
         );
     }
 
+    async fn run_imds_live_test(id: Option<UserAssignedId>) -> azure_core::Result<()> {
+        if std::env::var("IDENTITY_IMDS_AVAILABLE").is_err() {
+            println!("Skipped: IMDS isn't available");
+            return Ok(());
+        }
+
+        let credential = ManagedIdentityCredential::new(Some(ManagedIdentityCredentialOptions {
+            user_assigned_id: id,
+            ..Default::default()
+        }))
+        .expect("valid credential");
+
+        let token = credential.get_token(LIVE_TEST_SCOPES, None).await?;
+
+        assert!(!token.token.secret().is_empty());
+        assert_eq!(time::UtcOffset::UTC, token.expires_on.offset());
+        assert!(token.expires_on.unix_timestamp() > OffsetDateTime::now_utc().unix_timestamp());
+
+        Ok(())
+    }
+
     async fn run_imds_test(options: Option<ManagedIdentityCredentialOptions>) {
         let mut model = Request::new(
             "http://169.254.169.254/metadata/identity/oauth2/token"
@@ -408,11 +477,6 @@ mod tests {
         ).await;
     }
 
-    #[tokio::test]
-    async fn imds() {
-        run_imds_test(None).await;
-    }
-
     #[tokio::test]
     async fn imds_client_id() {
         run_imds_test(Some(ManagedIdentityCredentialOptions {
@@ -442,6 +506,16 @@ mod tests {
         .await;
     }
 
+    #[tokio::test]
+    async fn imds_system_assigned() {
+        run_imds_test(None).await;
+    }
+
+    #[recorded::test(live)]
+    async fn imds_system_assigned_live() -> azure_core::Result<()> {
+        run_imds_live_test(None).await
+    }
+
     #[tokio::test]
     async fn requires_one_scope() {
         let credential = ManagedIdentityCredential::new(None).expect("valid credential");
diff --git a/sdk/identity/azure_identity/tests/tools/deployed_live_test/Cargo.toml b/sdk/identity/azure_identity/tests/tools/deployed_live_test/Cargo.toml
@@ -0,0 +1,17 @@
+[package]
+name = "deployed_live_test"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[workspace]
+
+[dependencies]
+azure_identity = { path = "../../../" }
+azure_storage_blob = { path = "../../../../../storage/azure_storage_blob" }
+azure_core = { path = "../../../../../core/azure_core" }
+openssl = { version = "0.10", features = ["vendored"]}
+tokio = { version = "1.0", features = ["full"] }
+url = "2.5"
+axum = { version = "0.8", default-features = false, features = ["http1", "tokio", "query"] }
+serde = { version = "1.0", features = ["derive"] }
diff --git a/sdk/identity/azure_identity/tests/tools/deployed_live_test/src/main.rs b/sdk/identity/azure_identity/tests/tools/deployed_live_test/src/main.rs
@@ -0,0 +1,100 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use axum::{
+    extract::Query,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    routing::get,
+    Router,
+};
+use azure_core::credentials::TokenCredential;
+use azure_identity::{ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId};
+use azure_storage_blob::BlobServiceClient;
+use serde::Deserialize;
+use std::sync::Arc;
+
+#[derive(Deserialize)]
+#[serde(rename_all = "kebab-case")]
+struct Params {
+    client_id: Option<String>,
+    object_id: Option<String>,
+    resource_id: Option<String>,
+    storage_name: String,
+    test: String,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // cspell:ignore CUSTOMHANDLER
+    let port = std::env::var("FUNCTIONS_CUSTOMHANDLER_PORT").unwrap_or_else(|_| "8080".to_string());
+    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
+    println!("Listening on http://{}", listener.local_addr()?);
+
+    let router = Router::new().route("/api", get(handle_request));
+    axum::serve(listener, router)
+        .with_graceful_shutdown(async {
+            tokio::signal::ctrl_c()
+                .await
+                .expect("failed to handle ctrl-c");
+        })
+        .await?;
+
+    Ok(())
+}
+
+async fn handle_request(Query(params): Query<Params>) -> Response {
+    let credential = match params.test.as_str() {
+        "managed-identity" => {
+            let user_assigned_id = match (
+                params.client_id.as_ref(),
+                params.object_id.as_ref(),
+                params.resource_id.as_ref(),
+            ) {
+                (Some(id), None, None) => Some(UserAssignedId::ClientId(id.clone())),
+                (None, Some(id), None) => Some(UserAssignedId::ObjectId(id.clone())),
+                (None, None, Some(id)) => Some(UserAssignedId::ResourceId(id.clone())),
+                (None, None, None) => None,
+                _ => {
+                    return (
+                        StatusCode::BAD_REQUEST,
+                        "Multiple user-assigned identity parameters",
+                    )
+                        .into_response()
+                }
+            };
+            let options = ManagedIdentityCredentialOptions {
+                user_assigned_id,
+                ..Default::default()
+            };
+            match ManagedIdentityCredential::new(Some(options)) {
+                Ok(cred) => cred,
+                Err(e) => {
+                    return (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        format!("ManagedIdentityCredential::new returned '{e}'"),
+                    )
+                        .into_response()
+                }
+            }
+        }
+        test => return (StatusCode::BAD_REQUEST, format!("Unknown test '{test}'")).into_response(),
+    };
+
+    match try_storage(credential, &params.storage_name).await {
+        Ok(_) => StatusCode::OK.into_response(),
+        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(),
+    }
+}
+
+async fn try_storage(
+    credential: Arc<dyn TokenCredential>,
+    storage_name: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let endpoint = format!("https://{}.blob.core.windows.net", storage_name);
+    BlobServiceClient::new(endpoint.as_str(), credential, None)?
+        .get_properties(None)
+        .await
+        .map_err(|e| format!("BlobServiceClient::get_properties failed: {:?}", e).into())
+        .map(|_| ())
+}
diff --git a/sdk/identity/ci.yml b/sdk/identity/ci.yml
@@ -1,19 +1,27 @@
 # NOTE: Please refer to https://aka.ms/azsdk/engsys/ci-yaml before editing this file.
-
 trigger:
   branches:
     include:
-    - main
-    - hotfix/*
-    - release/*
+      - main
+      - hotfix/*
+      - release/*
   paths:
     include:
-    - sdk/identity/
+      - sdk/identity/
 
 extends:
   template: /eng/pipelines/templates/stages/archetype-sdk-client.yml
   parameters:
     ServiceDirectory: identity
     Artifacts:
-    - name: azure_identity
-      safeName: AzureIdentity
+      - name: azure_identity
+        safeName: AzureIdentity
+
+    ${{ if endsWith(variables['Build.DefinitionName'], 'weekly') }}:
+      RunLiveTests: true
+      PersistOidcToken: true
+      MatrixConfigs:
+        - Name: managed_identity_matrix
+          GenerateVMJobs: true
+          Path: sdk/identity/azure_identity/managed-identity-matrix.json
+          Selection: sparse
diff --git a/sdk/identity/test-resources-post.ps1 b/sdk/identity/test-resources-post.ps1
@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+# IMPORTANT: Do not invoke this file directly. Please instead run eng/common/TestResources/New-TestResources.ps1 from the repository root.
+
+param (
+  [hashtable] $AdditionalParameters = @{},
+  [hashtable] $DeploymentOutputs,
+
+  [Parameter(Mandatory = $true)]
+  [ValidateNotNullOrEmpty()]
+  [string] $SubscriptionId,
+
+  [Parameter(ParameterSetName = 'Provisioner', Mandatory = $true)]
+  [ValidateNotNullOrEmpty()]
+  [string] $TenantId,
+
+  [Parameter()]
+  [ValidatePattern('^[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}$')]
+  [string] $TestApplicationId,
+
+  [Parameter(Mandatory = $true)]
+  [ValidateNotNullOrEmpty()]
+  [string] $Environment,
+
+  # Captures any arguments from eng/New-TestResources.ps1 not declared here (no parameter errors).
+  [Parameter(ValueFromRemainingArguments = $true)]
+  $RemainingArguments
+)
+
+$ErrorActionPreference = 'Stop'
+$PSNativeCommandUseErrorActionPreference = $true
+
+if ($CI) {
+  if (!$AdditionalParameters['deployResources']) {
+    Write-Host "Skipping post-provisioning script because resources weren't deployed"
+    return
+  }
+  az cloud set -n $Environment
+  az login --federated-token $env:ARM_OIDC_TOKEN --service-principal -t $TenantId -u $TestApplicationId
+  az account set --subscription $SubscriptionId
+}
+
+Set-Location "$(git rev-parse --show-toplevel)/sdk/identity/azure_identity/tests/tools/deployed_live_test"
+
+Write-Host "##[group]Building test app"
+cargo install --path . --root target
+Write-Host "##[endgroup]"
+
+Write-Host "##[group]Building container image"
+az acr login -n $DeploymentOutputs['IDENTITY_ACR_NAME']
+$image = "$($DeploymentOutputs['IDENTITY_ACR_LOGIN_SERVER'])/live-test"
+Set-Content -Path Dockerfile -Value @"
+FROM mcr.microsoft.com/mirror/docker/library/ubuntu:24.04
+RUN apt update && apt install ca-certificates --no-install-recommends -y
+COPY target/bin/deployed_live_test .
+CMD ["./deployed_live_test"]
+"@
+docker build -t $image .
+docker push $image
+Write-Host "##[endgroup]"
+
+$rg = $DeploymentOutputs['IDENTITY_RESOURCE_GROUP']
+
+Write-Host "##[group]Deploying Azure Container Instance with user-assigned identity"
+$aciName = "azure-identity-test-user-assigned"
+az container create -g $rg -n $aciName --image $image `
+  --acr-identity $($DeploymentOutputs['IDENTITY_USER_ASSIGNED_IDENTITY']) `
+  --assign-identity $($DeploymentOutputs['IDENTITY_USER_ASSIGNED_IDENTITY']) `
+  --cpu 1 `
+  --ip-address Public `
+  --memory 1.0 `
+  --os-type Linux `
+  --ports 8080
+$aciIP = az container show -g $rg -n $aciName --query ipAddress.ip -o tsv
+Write-Host "##vso[task.setvariable variable=IDENTITY_ACI_IP_USER_ASSIGNED;]$aciIP"
+Write-Host "##[endgroup]"
diff --git a/sdk/identity/test-resources.bicep b/sdk/identity/test-resources.bicep
