[MSI v2] - Enable attestation in pop flows (#5496)

* attestation

* address pr comments

* remove console

* pr comments

* pr comments

Author: gladjohn

src/client/Microsoft.Identity.Client/ApiConfig/AcquireTokenForManagedIdentityParameterBuilder.cs

@@ -8,6 +8,7 @@
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.ApiConfig.Executors;
 using Microsoft.Identity.Client.ApiConfig.Parameters;
+using Microsoft.Identity.Client.ManagedIdentity;
 using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
 using Microsoft.Identity.Client.Utils;
 
@@ -80,6 +81,7 @@ public AcquireTokenForManagedIdentityParameterBuilder WithClaims(string claims)
         /// <inheritdoc/>
         internal override Task<AuthenticationResult> ExecuteInternalAsync(CancellationToken cancellationToken)
         {
+            ApplyMtlsPopAndAttestation(acquireTokenForManagedIdentityParameters: Parameters, acquireTokenCommonParameters: CommonParameters);
             return ManagedIdentityApplicationExecutor.ExecuteAsync(CommonParameters, Parameters, cancellationToken);
         }
 
@@ -93,5 +95,29 @@ internal override ApiEvent.ApiIds CalculateApiEventId()
 
             return ApiEvent.ApiIds.AcquireTokenForUserAssignedManagedIdentity;
         }
+
+        /// <summary>
+        /// TEST HOOK ONLY: Allows unit tests to inject a fake attestation-token provider
+        /// so we don't hit the real attestation service. Not part of the public API.
+        /// </summary>
+        internal AcquireTokenForManagedIdentityParameterBuilder WithAttestationProviderForTests(
+            Func<AttestationTokenInput, CancellationToken, Task<AttestationTokenResponse>> provider)
+        {
+            if (provider is null)
+            {
+                throw new ArgumentNullException(nameof(provider));
+            }
+
+            CommonParameters.AttestationTokenProvider = provider;
+            return this;
+        }
+
+        private static void ApplyMtlsPopAndAttestation(
+            AcquireTokenCommonParameters acquireTokenCommonParameters,
+            AcquireTokenForManagedIdentityParameters acquireTokenForManagedIdentityParameters)
+        {
+            acquireTokenForManagedIdentityParameters.IsMtlsPopRequested = acquireTokenCommonParameters.IsMtlsPopRequested;
+            acquireTokenForManagedIdentityParameters.AttestationTokenProvider ??= acquireTokenCommonParameters.AttestationTokenProvider;
+        }
     }
 }

src/client/Microsoft.Identity.Client/ApiConfig/Parameters/AcquireTokenForManagedIdentityParameters.cs

@@ -5,8 +5,10 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
+using Microsoft.Identity.Client.ManagedIdentity;
 
 namespace Microsoft.Identity.Client.ApiConfig.Parameters
 {
@@ -22,6 +24,8 @@ internal class AcquireTokenForManagedIdentityParameters : IAcquireTokenParameter
 
         public bool IsMtlsPopRequested { get; set; }
 
+        internal Func<AttestationTokenInput, CancellationToken, Task<AttestationTokenResponse>> AttestationTokenProvider { get; set; }
+
         public void LogParameters(ILoggerAdapter logger)
         {
             if (logger.IsLoggingEnabled(LogLevel.Info))

src/client/Microsoft.Identity.Client/AppConfig/ApplicationConfiguration.cs

@@ -17,6 +17,7 @@
 using Microsoft.Identity.Client.Internal.Broker;
 using Microsoft.Identity.Client.Internal.ClientCredential;
 using Microsoft.Identity.Client.Kerberos;
+using Microsoft.Identity.Client.ManagedIdentity;
 using Microsoft.Identity.Client.ManagedIdentity.V2;
 using Microsoft.Identity.Client.PlatformsCommon.Interfaces;
 using Microsoft.Identity.Client.UI;

src/client/Microsoft.Identity.Client/AppConfig/ManagedIdentityApplicationBuilder.cs

@@ -11,6 +11,7 @@
 using Microsoft.Identity.Client.AppConfig;
 using Microsoft.Identity.Client.Extensibility;
 using Microsoft.Identity.Client.Internal;
+using Microsoft.Identity.Client.ManagedIdentity;
 using Microsoft.Identity.Client.TelemetryCore;
 using Microsoft.Identity.Client.TelemetryCore.TelemetryClient;
 using Microsoft.Identity.Client.Utils;

src/client/Microsoft.Identity.Client/Internal/RequestContext.cs

@@ -5,8 +5,10 @@
 using System.Collections.Generic;
 using System.Security.Cryptography.X509Certificates;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Internal.Logger;
+using Microsoft.Identity.Client.ManagedIdentity;
 using Microsoft.Identity.Client.TelemetryCore;
 using Microsoft.Identity.Client.TelemetryCore.Internal.Events;
 using Microsoft.Identity.Client.TelemetryCore.TelemetryClient;
@@ -29,6 +31,8 @@ internal class RequestContext
 
         public X509Certificate2 MtlsCertificate { get; }
 
+        internal Func<AttestationTokenInput, CancellationToken, Task<AttestationTokenResponse>> AttestationTokenProvider { get; set; }
+
         public RequestContext(IServiceBundle serviceBundle, Guid correlationId, X509Certificate2 mtlsCertificate, CancellationToken cancellationToken = default)
         {
             ServiceBundle = serviceBundle ?? throw new ArgumentNullException(nameof(serviceBundle));

src/client/Microsoft.Identity.Client/Internal/Requests/ManagedIdentityAuthRequest.cs

@@ -200,6 +200,10 @@ private async Task<AuthenticationResult> SendTokenRequestForManagedIdentityAsync
 
             _managedIdentityParameters.IsMtlsPopRequested = AuthenticationRequestParameters.IsMtlsPopRequested;
 
+            // Ensure the attestation provider reaches RequestContext for IMDSv2
+            AuthenticationRequestParameters.RequestContext.AttestationTokenProvider ??=
+                _managedIdentityParameters.AttestationTokenProvider;
+
             ManagedIdentityResponse managedIdentityResponse =
                 await _managedIdentityClient
                 .SendTokenRequestForManagedIdentityAsync(AuthenticationRequestParameters.RequestContext, _managedIdentityParameters, cancellationToken)

src/client/Microsoft.Identity.Client/ManagedIdentity/V2/ImdsV2ManagedIdentitySource.cs

@@ -6,6 +6,7 @@
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.Core;
 using Microsoft.Identity.Client.Http;
@@ -31,7 +32,7 @@ public static async Task<CsrMetadata> GetCsrMetadataAsync(
             bool probeMode)
         {
 #if NET462
-            requestContext.Logger.Info(() => "[Managed Identity] IMDSv2 flow is not supported on .NET Framework 4.6.2. Cryptographic operations required for managed identity authentication are unavailable on this platform. Skipping IMDSv2 probe.");
+            requestContext.Logger.Info("[Managed Identity] IMDSv2 flow is not supported on .NET Framework 4.6.2. Cryptographic operations required for managed identity authentication are unavailable on this platform. Skipping IMDSv2 probe.");
             return await Task.FromResult<CsrMetadata>(null).ConfigureAwait(false);
 #else
             var queryParams = ImdsV2QueryParamsHelper(requestContext);
@@ -66,7 +67,7 @@ public static async Task<CsrMetadata> GetCsrMetadataAsync(
             {
                 if (probeMode)
                 {
-                    requestContext.Logger.Info(() => $"[Managed Identity] IMDSv2 CSR endpoint failure. Exception occurred while sending request to CSR metadata endpoint: ${ex}");
+                    requestContext.Logger.Info($"[Managed Identity] IMDSv2 CSR endpoint failure. Exception occurred while sending request to CSR metadata endpoint: {ex}");
                     return null;
                 }
                 else
@@ -187,7 +188,11 @@ internal ImdsV2ManagedIdentitySource(RequestContext requestContext) :
             base(requestContext, ManagedIdentitySource.ImdsV2)
         { }
 
-        private async Task<CertificateRequestResponse> ExecuteCertificateRequestAsync(string csr)
+        private async Task<CertificateRequestResponse> ExecuteCertificateRequestAsync(
+            string clientId,
+            string attestationEndpoint,
+            string csr,
+            ManagedIdentityKeyInfo managedIdentityKeyInfo)
         {
             var queryParams = ImdsV2QueryParamsHelper(_requestContext);
 
@@ -199,10 +204,32 @@ private async Task<CertificateRequestResponse> ExecuteCertificateRequestAsync(st
                 { OAuth2Header.XMsCorrelationId, _requestContext.CorrelationId.ToString() }
             };
 
+            if (_isMtlsPopRequested && managedIdentityKeyInfo.Type != ManagedIdentityKeyType.KeyGuard)
+            {
+                throw new MsalClientException(
+                    "mtls_pop_requires_keyguard",
+                    "[ImdsV2] mTLS Proof-of-Possession requires a KeyGuard-backed key. Enable KeyGuard or use a KeyGuard-supported environment.");
+            }
+
+            // TODO: : Normalize and validate attestation endpoint Code needs to be removed 
+            // once IMDS team start returning full URI
+            Uri normalizedEndpoint = NormalizeAttestationEndpoint(attestationEndpoint, _requestContext.Logger);
+
+            // Ask helper for JWT only for KeyGuard keys
+            string attestationJwt = string.Empty;
+            if (managedIdentityKeyInfo.Type == ManagedIdentityKeyType.KeyGuard)
+            {
+                attestationJwt = await GetAttestationJwtAsync(
+                    clientId,
+                    normalizedEndpoint,
+                    managedIdentityKeyInfo,
+                    _requestContext.UserCancellationToken).ConfigureAwait(false);
+            }
+
             var certificateRequestBody = new CertificateRequestBody()
             {
                 Csr = csr,
-                // AttestationToken = "fake_attestation_token" TODO: implement attestation token
+                AttestationToken = attestationJwt
             };
 
             string body = JsonHelper.SerializeToJson(certificateRequestBody);
@@ -257,12 +284,21 @@ protected override async Task<ManagedIdentityRequest> CreateRequestAsync(string
         {
             var csrMetadata = await GetCsrMetadataAsync(_requestContext, false).ConfigureAwait(false);
 
-            var keyInfo = await _requestContext.ServiceBundle.PlatformProxy.ManagedIdentityKeyProvider
-                .GetOrCreateKeyAsync(_requestContext.Logger, _requestContext.UserCancellationToken).ConfigureAwait(false);
+            IManagedIdentityKeyProvider keyProvider = _requestContext.ServiceBundle.PlatformProxy.ManagedIdentityKeyProvider;
+
+            ManagedIdentityKeyInfo keyInfo = await keyProvider
+                .GetOrCreateKeyAsync(
+                _requestContext.Logger, 
+                _requestContext.UserCancellationToken)
+                .ConfigureAwait(false);
 
             var (csr, privateKey) = _requestContext.ServiceBundle.Config.CsrFactory.Generate(keyInfo.Key, csrMetadata.ClientId, csrMetadata.TenantId, csrMetadata.CuId);
 
-            var certificateRequestResponse = await ExecuteCertificateRequestAsync(csr).ConfigureAwait(false);
+            var certificateRequestResponse = await ExecuteCertificateRequestAsync(
+                csrMetadata.ClientId,
+                csrMetadata.AttestationEndpoint,
+                csr,
+                keyInfo).ConfigureAwait(false);
 
             // transform certificateRequestResponse.Certificate to x509 with private key
             var mtlsCertificate = CommonCryptographyManager.AttachPrivateKeyToCert(
@@ -302,12 +338,117 @@ private static string ImdsV2QueryParamsHelper(RequestContext requestContext)
                 requestContext.ServiceBundle.Config.ManagedIdentityId.IdType,
                 requestContext.ServiceBundle.Config.ManagedIdentityId.UserAssignedId,
                 requestContext.Logger);
+            
             if (userAssignedIdQueryParam != null)
             {
                 queryParams += $"&{userAssignedIdQueryParam.Value.Key}={userAssignedIdQueryParam.Value.Value}";
             }
 
             return queryParams;
         }
+
+        /// <summary>
+        /// Obtains an attestation JWT for the KeyGuard/CSR payload using the configured
+        /// attestation provider and normalized endpoint.
+        /// </summary>
+        /// <param name="clientId">Client ID to be sent to the attestation provider.</param>
+        /// <param name="attestationEndpoint">The attestation endpoint.</param>
+        /// <param name="keyInfo">The key information.</param>
+        /// <param name="cancellationToken">Cancellation token.</param>
+        /// <returns>JWT string suitable for the IMDSv2 attested POP flow.</returns>
+        /// <exception cref="MsalClientException">Wraps client/network failures.</exception>
+
+        private async Task<string> GetAttestationJwtAsync(
+            string clientId, 
+            Uri attestationEndpoint, 
+            ManagedIdentityKeyInfo keyInfo, 
+            CancellationToken cancellationToken)
+        {
+            // Provider is a local dependency; missing provider is a client error
+            var provider = _requestContext.AttestationTokenProvider;
+
+            // KeyGuard requires RSACng on Windows
+            if (keyInfo.Type == ManagedIdentityKeyType.KeyGuard &&
+                keyInfo.Key is not System.Security.Cryptography.RSACng rsaCng)
+            {
+                throw new MsalClientException(
+                    "keyguard_requires_cng",
+                    "[ImdsV2] KeyGuard attestation currently supports only RSA CNG keys on Windows.");
+            }
+
+            // Attestation token input
+            var input = new AttestationTokenInput
+            {
+                ClientId = clientId,
+                AttestationEndpoint = attestationEndpoint,
+                KeyHandle = (keyInfo.Key as System.Security.Cryptography.RSACng)?.Key.Handle
+            };
+
+            // response from provider 
+            var response = await provider(input, cancellationToken).ConfigureAwait(false);
+
+            // Validate response
+            if (response == null || string.IsNullOrWhiteSpace(response.AttestationToken))
+            {
+                throw new MsalClientException(
+                    "attestation_failed",
+                    "[ImdsV2] Attestation provider failed to return an attestation token.");
+            }
+
+            // Return the JWT
+            return response.AttestationToken;
+        }
+
+        //To-do : Remove this method once IMDS team start returning full URI
+        /// <summary>
+        /// Temporarily normalize attestation endpoint values to a full https:// URI.
+        /// IMDS team will eventually return a full URI. 
+        /// </summary>
+        /// <param name="rawEndpoint"></param>
+        /// <param name="logger"></param>
+        /// <returns></returns>
+        private static Uri NormalizeAttestationEndpoint(string rawEndpoint, ILoggerAdapter logger)
+        {
+            if (string.IsNullOrWhiteSpace(rawEndpoint))
+            {
+                return null;
+            }
+
+            // Trim whitespace
+            rawEndpoint = rawEndpoint.Trim();
+
+            // If it already parses as an absolute URI with https, keep it.
+            if (Uri.TryCreate(rawEndpoint, UriKind.Absolute, out var absolute) &&
+                (absolute.Scheme.Equals("https", StringComparison.OrdinalIgnoreCase)))
+            {
+                return absolute;
+            }
+
+            // If it has no scheme (common service behavior returning only host)
+            // prepend https:// and try again.
+            if (!rawEndpoint.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
+            {
+                var candidate = "https://" + rawEndpoint;
+                if (Uri.TryCreate(candidate, UriKind.Absolute, out var httpsUri))
+                {
+                    logger.Info(() => $"[Managed Identity] Normalized attestation endpoint '{rawEndpoint}' -> '{httpsUri.ToString()}'.");
+                    return httpsUri;
+                }
+            }
+
+            // Final attempt: reject http (non‑TLS) or malformed
+            if (Uri.TryCreate(rawEndpoint, UriKind.Absolute, out var anyUri))
+            {
+                if (!anyUri.Scheme.Equals("https", StringComparison.OrdinalIgnoreCase))
+                {
+                    logger.Warning($"[Managed Identity] Attestation endpoint uses unsupported scheme '{anyUri.Scheme}'. HTTPS is required.");
+                    return null;
+                }
+                return anyUri;
+            }
+
+            logger.Warning($"[Managed Identity] Failed to normalize attestation endpoint value '{rawEndpoint}'.");
+            return null;
+        }
     }
 }