From 0eb1e6764b676c30e97b93c94ea83221f59d9155 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ben=20Carpenter=20=F0=9F=9B=AB?=
 <carpenterben@microsoft.com>
Date: Wed, 17 Jul 2024 15:01:30 -0700
Subject: [PATCH] add managed identity option for authorization

---
 .../Utility/AuthorizationHandler.cs           | 79 +++++++++++++++----
 .../AuthorizationHandlerCredentials.cs        |  8 +-
 2 files changed, 72 insertions(+), 15 deletions(-)

diff --git a/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandler.cs b/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandler.cs
index 2950ad3..bd6c41f 100644
--- a/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandler.cs
+++ b/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandler.cs
@@ -125,27 +125,78 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
         return response;
     }
 
-    private async Task<bool> ObtainAccessToken()
+    private async Task<AccessToken> ObtainAccessToken(CancellationToken cancellationToken = default)
     {
-        bool IsSuccess = false;
-        string DevCenterTokenUrl = string.Format("https://login.microsoftonline.com/{0}/oauth2/token", _authCredentials.TenantId);
+        if (!string.IsNullOrEmpty(_authCredentials.ManagedIdentityClientId))
+        {
+            return await GetClientAssertionTokenAsync(cancellationToken);
+        }
+        else
+        {
+            return await GetTokenUsingClientSecretAsync(cancellationToken);
+        }
+    }
+
+    private async Task<AccessToken> GetClientAssertionTokenAsync(CancellationToken cancellationToken)
+    {
+        AccessToken token = default;
 
-        using (HttpClient client = new())
+        var clientAssertionCredential = new ClientAssertionCredential(
+            _authCredentials.TenantId,
+            _authCredentials.ClientId,
+            async (token) => await GetTokenUsingManagedIdentityAsync(cancellationToken)
+        );
+
+        token = await clientAssertionCredential.GetTokenAsync(
+            new TokenRequestContext(new[] { "https://manage.devcenter.microsoft.com/.default" }),
+            cancellationToken
+        );
+
+        if (!string.IsNullOrEmpty(token.Token))
         {
-            client.Timeout = _httpTimeout;
-            Uri restApi = new(DevCenterTokenUrl);
+            _accessToken = token.Token;
+        }
 
-            ClientSecretCredential credential = new(_authCredentials.TenantId, _authCredentials.ClientId, _authCredentials.Key);
-            AccessToken token = await credential.GetTokenAsync(new TokenRequestContext(scopes: new string[] { "https://manage.devcenter.microsoft.com/.default" }));
+        return token;
+    }
 
-            if (string.IsNullOrEmpty(token.Token) == false)
-            {
-                _accessToken = token.Token;
-                IsSuccess = true;
-            }
+    /// <summary>
+    /// Callback function for <see cref="GetClientAssertionTokenAsync"/>
+    /// </summary>
+    /// <param name="cancellationToken"></param>
+    /// <returns></returns>
+    private async Task<string> GetTokenUsingManagedIdentityAsync(CancellationToken cancellationToken)
+    {
+        var credential = new ManagedIdentityCredential(_authCredentials.ManagedIdentityClientId);
+
+        var token = await credential.GetTokenAsync(
+            new TokenRequestContext(new[] { _authCredentials.Scope }),
+            cancellationToken
+        );
+
+        return token.Token;
+    }
+
+    private async Task<AccessToken> GetTokenUsingClientSecretAsync(CancellationToken cancellationToken)
+    {
+
+        var credential = new ClientSecretCredential(
+            _authCredentials.TenantId,
+            _authCredentials.ClientId,
+            _authCredentials.Key
+        );
+
+        var token = await credential.GetTokenAsync(
+            new TokenRequestContext(new[] { "https://manage.devcenter.microsoft.com/.default" }),
+            cancellationToken
+        );
+
+        if (!string.IsNullOrEmpty(token.Token))
+        {
+            _accessToken = token.Token;
         }
 
-        return IsSuccess;
+        return token;
     }
 
     //
diff --git a/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandlerCredentials.cs b/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandlerCredentials.cs
index aa6cfa3..3a520a9 100644
--- a/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandlerCredentials.cs
+++ b/src/Microsoft.Devices.HardwareDevCenterManager/Utility/AuthorizationHandlerCredentials.cs
@@ -1,7 +1,7 @@
 /*++
     Copyright (c) Microsoft Corporation. All rights reserved.
 
-    Licensed under the MIT license. See LICENSE file in the project root for full license information.  
+    Licensed under the MIT license. See LICENSE file in the project root for full license information.
 --*/
 
 using System.Text.Json.Serialization;
@@ -24,4 +24,10 @@ public class AuthorizationHandlerCredentials
 
     [JsonPropertyName("urlPrefix")]
     public System.Uri UrlPrefix { get; set; }
+
+    [JsonPropertyName("managedIdentityClientId")]
+    public string ManagedIdentityClientId { get; set; }
+
+    [JsonPropertyName("scope")]
+    public string Scope { get; set; }
 }