ImdsV2ManagedIdentitySource.cs

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Identity.Client.Core;
using Microsoft.Identity.Client.Http;
using Microsoft.Identity.Client.Http.Retry;
using Microsoft.Identity.Client.Internal;
using Microsoft.Identity.Client.OAuth2;
using Microsoft.Identity.Client.OAuth2.Throttling;
using Microsoft.Identity.Client.PlatformsCommon.Shared;
using Microsoft.Identity.Client.Utils;

namespace Microsoft.Identity.Client.ManagedIdentity.V2
{
    internal class ImdsV2ManagedIdentitySource : IImdsV2MtlsBindingSource
    {
        // Central, process-local cache for mTLS binding (cert + endpoint + canonical client_id).
        internal static readonly ICertificateCache s_mtlsCertificateCache = new InMemoryCertificateCache();

        private readonly RequestContext _requestContext;
        private readonly IMtlsCertificateCache _mtlsCache;
        private bool _isMtlsPopRequested;
        private bool _isMtlsBearerRequested;
        private Func<string, SafeHandle, string, string, ILoggerAdapter, CancellationToken, Task<string>> _attestationTokenProvider;

        // used in unit tests
        public const string ApiVersionQueryParam = "cred-api-version";
        public const string ImdsV2ApiVersion = "2.0";
        public const string CsrMetadataPath = "/metadata/identity/getplatformmetadata";
        public const string CertificateRequestPath = "/metadata/identity/issuecredential";
        public const string AcquireEntraTokenPath = "/oauth2/v2.0/token";
        private const string AttestationTagEnabled = "#att=1";
        private const string AttestationTagDisabled = "#att=0";

        public static async Task<CsrMetadata> GetCsrMetadataAsync(RequestContext requestContext)
        {
            var queryParams = ImdsManagedIdentitySource.ImdsQueryParamsHelper(requestContext, ApiVersionQueryParam, ImdsV2ApiVersion);

            var headers = new Dictionary<string, string>
            {
                { "Metadata", "true" },
                { OAuth2Header.XMsCorrelationId, requestContext.CorrelationId.ToString() }
            };

            IRetryPolicyFactory retryPolicyFactory = requestContext.ServiceBundle.Config.RetryPolicyFactory;
            IRetryPolicy retryPolicy = retryPolicyFactory.GetRetryPolicy(RequestType.Imds);

            HttpResponse response = null;

            try
            {
                response = await requestContext.ServiceBundle.HttpManager.SendRequestAsync(
                    ImdsManagedIdentitySource.GetValidatedEndpoint(requestContext.Logger, CsrMetadataPath, queryParams),
                    headers,
                    body: null,
                    method: HttpMethod.Get,
                    logger: requestContext.Logger,
                    doNotThrow: false,
                    mtlsCertificate: null,
                    validateServerCertificate: null,
                    cancellationToken: requestContext.UserCancellationToken,
                    retryPolicy: retryPolicy)
                .ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                ThrowCsrMetadataRequestException(
                    "ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed.",
                    ex);
            }

            if (response.StatusCode != HttpStatusCode.OK)
            {
                ThrowCsrMetadataRequestException(
                    $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed due to HTTP error. Status code: {response.StatusCode} Body: {response.Body}",
                    null,
                    (int)response.StatusCode);
            }

            if (!ValidateCsrMetadataResponse(response, requestContext.Logger))
            {
                return null;
            }

            return TryCreateCsrMetadata(response, requestContext.Logger);
        }

        private static void ThrowCsrMetadataRequestException(
            String errorMessage,
            Exception ex = null,
            int? statusCode = null)
        {
            // A 404 from the IMDSv2 CSR endpoint indicates that the host supports only IMDSv1.
            // This happens when WithMtlsProofOfPossession() is used without a prior
            // GetManagedIdentitySourceAsync() call: MSAL routes directly to IMDSv2, and
            // on an IMDSv1-only host the /getplatformmetadata endpoint does not exist.
            // Translate to a client error so callers know mTLS PoP is not supported here.
            if (statusCode == (int)HttpStatusCode.NotFound)
            {
                throw new MsalClientException(
                    MsalError.MtlsPopTokenNotSupportedinImdsV1,
                    MsalErrorMessage.MtlsPopTokenNotSupportedinImdsV1);
            }

            throw MsalServiceExceptionFactory.CreateManagedIdentityException(
                MsalError.ManagedIdentityRequestFailed,
                $"[ImdsV2] {errorMessage}",
                ex,
                ManagedIdentitySource.Imds,
                statusCode);
        }

        private static bool ValidateCsrMetadataResponse(
            HttpResponse response,
            ILoggerAdapter logger)
        {
            string serverHeader = response.HeadersAsDictionary
                .FirstOrDefault((kvp) => {
                    return string.Equals(kvp.Key, "server", StringComparison.OrdinalIgnoreCase);
                }).Value;

            if (serverHeader == null)
            {
                ThrowCsrMetadataRequestException(
                    $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because response doesn't have server header. Status code: {response.StatusCode} Body: {response.Body}",
                    null,
                    (int)response.StatusCode);
            }

            if (!serverHeader.Contains("IMDS", StringComparison.OrdinalIgnoreCase))
            {
                ThrowCsrMetadataRequestException(
                    $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because the 'server' header format is invalid. Extracted server header: {serverHeader}. Status code: {response.StatusCode} Body: {response.Body}",
                    null,
                    (int)response.StatusCode);
            }

            return true;
        }

        private static CsrMetadata TryCreateCsrMetadata(
            HttpResponse response,
            ILoggerAdapter logger)
        {
            CsrMetadata csrMetadata = JsonHelper.DeserializeFromJson<CsrMetadata>(response.Body);
            if (!CsrMetadata.ValidateCsrMetadata(csrMetadata))
            {
                ThrowCsrMetadataRequestException(
                    $"ImdsV2ManagedIdentitySource.GetCsrMetadataAsync failed because the CsrMetadata response is invalid. Status code: {response.StatusCode} Body: {response.Body}",
                    null,
                    (int)response.StatusCode);
            }

            logger.Info(() => "[Managed Identity] IMDSv2 managed identity is available.");
            return csrMetadata;
        }

        public static ImdsV2ManagedIdentitySource Create(RequestContext requestContext)
        {
            return new ImdsV2ManagedIdentitySource(requestContext);
        }

        internal ImdsV2ManagedIdentitySource(RequestContext requestContext)
            : this(requestContext,
                  new MtlsBindingCache(s_mtlsCertificateCache, PersistentCertificateCacheFactory
                      .Create(requestContext.Logger)))
        {
        }

        internal ImdsV2ManagedIdentitySource(
            RequestContext requestContext,
            IMtlsCertificateCache mtlsCache)
        {
            _requestContext = requestContext;
            _mtlsCache = mtlsCache ?? throw new ArgumentNullException(nameof(mtlsCache));
        }

        /// <summary>
        /// Detects if the exception was caused by a SCHANNEL failure during mTLS authentication, 
        /// which can occur if the client certificate becomes invalid.
        /// </summary>
        /// <param name="ex"></param>
        /// <returns></returns>
        internal static bool IsSchanelFailure(MsalServiceException ex)
        {
            for (Exception e = ex; e != null; e = e.InnerException)
            {
                if (e is System.Net.Sockets.SocketException se &&
                    (se.ErrorCode == 10054 || se.SocketErrorCode == System.Net.Sockets.SocketError.ConnectionReset))
                {
                    return true;
                }

                if (e is System.Security.Authentication.AuthenticationException)
                {
                    return true;
                }
            }

            return false;
        }

        private async Task<CertificateRequestResponse> ExecuteCertificateRequestAsync(
            string clientId,
            string attestationEndpoint,
            string csr,
            ManagedIdentityKeyInfo managedIdentityKeyInfo)
        {
            var queryParams = ImdsManagedIdentitySource.ImdsQueryParamsHelper(_requestContext, ApiVersionQueryParam, ImdsV2ApiVersion);

            // TODO: add bypass_cache query param in case of token revocation. Boolean: true/false

            var headers = new Dictionary<string, string>
            {
                { "Metadata", "true" },
                { OAuth2Header.XMsCorrelationId, _requestContext.CorrelationId.ToString() }
            };

            // Attempt attestation only for KeyGuard keys when provider is available
            // For non-KeyGuard keys (Hardware, InMemory), proceed with non-attested flow
            string attestationJwt = string.Empty;
            var attestationUri = new Uri(attestationEndpoint);

            if (managedIdentityKeyInfo.Type == ManagedIdentityKeyType.KeyGuard)
            {
                attestationJwt = await GetAttestationJwtAsync(
                    clientId,
                    attestationUri,
                    managedIdentityKeyInfo,
                    _requestContext.UserCancellationToken).ConfigureAwait(false);
            }
            else
            {
                _requestContext.Logger.Info($"[ImdsV2] Using {managedIdentityKeyInfo.Type} key. Proceeding with non-attested mTLS PoP flow.");
            }

            var certificateRequestBody = new CertificateRequestBody()
            {
                Csr = csr,
                AttestationToken = attestationJwt
            };

            string body = JsonHelper.SerializeToJson(certificateRequestBody);

            IRetryPolicyFactory retryPolicyFactory = _requestContext.ServiceBundle.Config.RetryPolicyFactory;
            IRetryPolicy retryPolicy = retryPolicyFactory.GetRetryPolicy(RequestType.Imds);

            HttpResponse response = null;

            try
            {
                response = await _requestContext.ServiceBundle.HttpManager.SendRequestAsync(
                    ImdsManagedIdentitySource.GetValidatedEndpoint(_requestContext.Logger, CertificateRequestPath, queryParams),
                    headers,
                    body: new StringContent(body, System.Text.Encoding.UTF8, "application/json"),
                    method: HttpMethod.Post,
                    logger: _requestContext.Logger,
                    doNotThrow: false,
                    mtlsCertificate: null,
                    validateServerCertificate: null,
                    cancellationToken: _requestContext.UserCancellationToken,
                    retryPolicy: retryPolicy)
                .ConfigureAwait(false);
            }
            catch (Exception ex)
            {
                int? statusCode = response != null ? (int?)response.StatusCode : null;

                throw MsalServiceExceptionFactory.CreateManagedIdentityException(
                    MsalError.ManagedIdentityRequestFailed,
                    "[ImdsV2] ImdsV2ManagedIdentitySource.ExecuteCertificateRequestAsync failed.",
                    ex,
                    ManagedIdentitySource.Imds,
                    statusCode);
            }

            if (response.StatusCode != HttpStatusCode.OK)
            {
                throw MsalServiceExceptionFactory.CreateManagedIdentityException(
                    MsalError.ManagedIdentityRequestFailed,
                    $"[ImdsV2] ImdsV2ManagedIdentitySource.ExecuteCertificateRequestAsync failed due to HTTP error. Status code: {response.StatusCode} Body: {response.Body}",
                    null,
                    ManagedIdentitySource.Imds,
                    (int)response.StatusCode);
            }

            var certificateRequestResponse = JsonHelper.DeserializeFromJson<CertificateRequestResponse>(response.Body);
            CertificateRequestResponse.Validate(certificateRequestResponse);

            return certificateRequestResponse;
        }

        /// <summary>
        /// Performs the cert-mint flow (/getplatformmetadata + /issuecredential) and returns the
        /// resulting mTLS binding (cert + ESTS-R endpoint + canonical client_id). Extracted so the
        /// binding can be reused by the internal-exchange delegation path without building the
        /// bespoke token request.
        /// </summary>
        private async Task<MtlsBindingInfo> AcquireMtlsBindingAsync()
        {
            CsrMetadata csrMetadata = await GetCsrMetadataAsync(_requestContext).ConfigureAwait(false);

            // Early validation: Fail-fast if KeyGuard is required (mTLS PoP or mTLS Bearer) but unavailable.
            // This check happens before any network calls to avoid wasted round-trips.
            // Note: This creates/retrieves the key, but on cache hit scenarios (below),
            // this may be the only key access needed.
            if (_isMtlsPopRequested || _isMtlsBearerRequested)
            {
                IManagedIdentityKeyProvider keyProvider = _requestContext.ServiceBundle.PlatformProxy.ManagedIdentityKeyProvider;
                ManagedIdentityKeyInfo keyInfo = await keyProvider
                    .GetOrCreateKeyAsync(_requestContext.Logger, _requestContext.UserCancellationToken)
                    .ConfigureAwait(false);

                if (keyInfo.Type != ManagedIdentityKeyType.KeyGuard)
                {
                    string flowName = _isMtlsPopRequested ? "mTLS Proof-of-Possession" : "mTLS Bearer";
                    throw new MsalClientException(
                        "credential_guard_not_available",
                        $"[ImdsV2] {flowName} currently requires a KeyGuard key, but this host produced a '{keyInfo.Type}' key. " +
                        "The host may report Software-strength binding capability (which means it can bind a token to a key), " +
                        "but the IMDSv2 attested flow only accepts VBS-isolated KeyGuard keys today. " +
                        "Ensure Virtualization-based Security (VBS)/KeyGuard is enabled on the host.");
                }
            }

            string certCacheKey = GetMtlsCertCacheKey();

            // Get or create mTLS binding (cert + endpoint + client_id) from cache.
            // The factory delegate only executes on cache miss.
            MtlsBindingInfo mtlsBinding = await GetOrCreateMtlsBindingAsync(
                cacheKey: certCacheKey,
                async () =>  // Factory: only invoked if cert is not in cache
                {
                    IManagedIdentityKeyProvider keyProvider = _requestContext.ServiceBundle.PlatformProxy.ManagedIdentityKeyProvider;

                    // Second GetOrCreateKeyAsync call: Required for CSR generation on cache miss.
                    // If the cert is cached, this entire factory delegate is skipped.
                    // If validation above succeeded, this call returns the same cached key immediately.
                    ManagedIdentityKeyInfo keyInfo = await keyProvider
                        .GetOrCreateKeyAsync(_requestContext.Logger, _requestContext.UserCancellationToken)
                        .ConfigureAwait(false);

                    var csrAndKey = _requestContext.ServiceBundle.Config.CsrFactory.Generate(
                        keyInfo.Key,
                        csrMetadata.ClientId,
                        csrMetadata.TenantId,
                        csrMetadata.CuId);

                    string csr = csrAndKey.csrPem;
                    var privateKey = csrAndKey.privateKey;

                    var certificateRequestResponse = await ExecuteCertificateRequestAsync(
                        csrMetadata.ClientId,
                        csrMetadata.AttestationEndpoint,
                        csr,
                        keyInfo).ConfigureAwait(false);

                    X509Certificate2 mtlsCertificate = CommonCryptographyManager.AttachPrivateKeyToCert(
                        certificateRequestResponse.Certificate,
                        privateKey);

                    // Base endpoint = "{mtlsAuthEndpoint}/{tenantId}"
                    string endpointBase =
                        (certificateRequestResponse.MtlsAuthenticationEndpoint).TrimEnd('/') +
                        "/" +
                        (certificateRequestResponse.TenantId).Trim('/');

                    // Canonical GUID to use as client_id in the token call
                    string clientIdGuid = certificateRequestResponse.ClientId;

                    return new MtlsBindingInfo(mtlsCertificate, endpointBase, clientIdGuid);

                },
                _requestContext.UserCancellationToken,
                _requestContext.Logger)
                .ConfigureAwait(false);

            return mtlsBinding;
        }

        /// <summary>
        /// Mint-only entrypoint used by the internal-exchange delegation path. Sets the attestation
        /// provider and mTLS-PoP flag from the request parameters, optionally evicts a rejected cert
        /// (invalid_client / SCHANNEL re-mint), and returns the mTLS binding. Does NOT send the token request.
        /// </summary>
        /// <remarks>
        /// IMDSv2 delegates its token leg to MSAL's internal TokenClient exchange (see
        /// ManagedIdentityAuthRequest.SendDelegatedImdsV2TokenRequestAsync). DO NOT restore a bespoke
        /// token request: the IMDSv2 token leg must go through TokenClient so client-originated claims,
        /// client-capability (CP1) merge, and claims-based cache keying are preserved.
        /// </remarks>
        public async Task<MtlsBindingInfo> AcquireMtlsBindingForDelegationAsync(
            ApiConfig.Parameters.AcquireTokenForManagedIdentityParameters parameters,
            bool forceRemint,
            CancellationToken cancellationToken)
        {
            _attestationTokenProvider = parameters.AttestationTokenProvider;
            _isMtlsPopRequested = parameters.IsMtlsPopRequested;
            _isMtlsBearerRequested = parameters.IsMtlsBearerRequested;

            // Ensure at least one IMDSv2 attested flag is set; default to PoP for backward compatibility
            // with callers that do not set either flag explicitly.
            if (!_isMtlsPopRequested && !_isMtlsBearerRequested)
            {
                _isMtlsPopRequested = true;
            }

            if (forceRemint && _mtlsCache is MtlsBindingCache bindingCache)
            {
                bindingCache.RemoveBadCert(GetMtlsCertCacheKey(), _requestContext.Logger);
            }

            cancellationToken.ThrowIfCancellationRequested();

            return await AcquireMtlsBindingAsync().ConfigureAwait(false);
        }

        /// <summary>
        /// Obtains an attestation JWT for the Credential Guard/CSR payload using the configured
        /// attestation token provider delegate.
        /// </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, or null for non-attested flow.</returns>
        private async Task<string> GetAttestationJwtAsync(
            string clientId,
            Uri attestationEndpoint,
            ManagedIdentityKeyInfo keyInfo,
            CancellationToken cancellationToken)
        {
            // Check if attestation token provider has been configured
            if (_attestationTokenProvider == null)
            {
                _requestContext.Logger.Info("[ImdsV2] Attestation token provider not configured. Proceeding with non-attested flow.");
                return null; // Null attestation token indicates non-attested flow
            }

            // Credential Guard requires RSACng on Windows
            if (keyInfo.Key is not System.Security.Cryptography.RSACng rsaCng)
            {
                throw new MsalClientException(
                    "credential_guard_requires_cng",
                    "[ImdsV2] Credential Guard attestation currently supports only RSA CNG keys on Windows.");
            }

            try
            {
                // Call the attestation token provider delegate.
                // Prefer the CNG key name (stable for persisted/KSP keys).
                // For ephemeral keys (no name), derive a stable identifier from the public key
                // fingerprint so that the same key handle maps to the same cache entry while
                // distinct ephemeral keys get distinct entries.
                string keyId = rsaCng.Key.KeyName ?? GetPublicKeyFingerprint(rsaCng);
                string attestationJwt = await _attestationTokenProvider(
                    attestationEndpoint.AbsoluteUri,
                    rsaCng.Key.Handle,
                    clientId,
                    keyId,
                    _requestContext.Logger,
                    cancellationToken).ConfigureAwait(false);

                if (string.IsNullOrWhiteSpace(attestationJwt))
                {
                    _requestContext.Logger.Info("[ImdsV2] Attestation provider returned null/empty JWT. Proceeding with non-attested flow.");
                    return null;
                }

                return attestationJwt;
            }
            catch (Exception ex)
            {
                throw new MsalClientException(
                    "attestation_failed",
                    $"[ImdsV2] Attestation token provider failed: {ex.Message}",
                    ex);
            }
        }

        private Task<MtlsBindingInfo> GetOrCreateMtlsBindingAsync(
            string cacheKey,
            Func<Task<MtlsBindingInfo>> factory,
            CancellationToken cancellationToken,
            ILoggerAdapter logger)
        {
            return _mtlsCache.GetOrCreateAsync(cacheKey, factory, cancellationToken, logger);
        }

        private string GetMtlsCertCacheKey()
        {
            // Today you use Config.ClientId as the base alias. Keep that unchanged.
            // Just disambiguate by whether WithAttestationSupport() was configured.
            string baseKey = _requestContext.ServiceBundle.Config.ClientId;

            // FriendlyName encoder forbids '|', CR/LF, NULL. "#att=*" is safe.
            return baseKey + (_attestationTokenProvider != null ? AttestationTagEnabled : AttestationTagDisabled);
        }

        internal static void ResetCertCacheForTest()
        {
            // Clear caches so each test starts fresh
            if (s_mtlsCertificateCache != null)
            {
                s_mtlsCertificateCache.Clear();
            }
        }

        /// <summary>
        /// Computes a stable hex fingerprint of the RSA public key.
        /// Used as a cache key for ephemeral CNG keys that have no key name.
        /// Compatible with .NET Framework 4.6.2 and netstandard2.0.
        /// </summary>
        private static string GetPublicKeyFingerprint(RSA rsa)
        {
            RSAParameters p = rsa.ExportParameters(includePrivateParameters: false);
            // Concatenate Modulus + Exponent as a stable, unique representation of the public key.
            byte[] combined = new byte[p.Modulus.Length + p.Exponent.Length];
            Buffer.BlockCopy(p.Modulus, 0, combined, 0, p.Modulus.Length);
            Buffer.BlockCopy(p.Exponent, 0, combined, p.Modulus.Length, p.Exponent.Length);
            using var sha256 = SHA256.Create();
            byte[] hash = sha256.ComputeHash(combined);
            return BitConverter.ToString(hash).Replace("-", string.Empty);
        }
    }
}