Add code and unit tests for cloud shell

neha-bhargava · neha-bhargava · commit 0913e044d60d · 2023-09-18T11:00:12.000-07:00
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractManagedIdentitySource.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/AbstractManagedIdentitySource.java
@@ -3,19 +3,12 @@
 
 package com.microsoft.aad.msal4j;
 
-import com.nimbusds.oauth2.sdk.ParseException;
-import com.nimbusds.oauth2.sdk.SerializeException;
-import com.nimbusds.oauth2.sdk.http.HTTPRequest;
-import com.nimbusds.oauth2.sdk.http.HTTPResponse;
 import lombok.Getter;
 import lombok.Setter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.beans.Encoder;
-import java.io.IOException;
 import java.net.HttpURLConnection;
-import java.net.MalformedURLException;
 import java.net.SocketException;
 import java.net.URISyntaxException;
 
@@ -56,7 +49,15 @@ public ManagedIdentityResponse getManagedIdentityResponse(
         IHttpResponse response;
 
         try {
-            HttpRequest httpRequest = new HttpRequest(HttpMethod.GET, managedIdentityRequest.computeURI().toString(), managedIdentityRequest.headers);
+
+            HttpRequest httpRequest = managedIdentityRequest.method.equals(HttpMethod.GET) ?
+                    new HttpRequest(HttpMethod.GET,
+                            managedIdentityRequest.computeURI().toString(),
+                            managedIdentityRequest.headers) :
+                    new HttpRequest(HttpMethod.POST,
+                            managedIdentityRequest.computeURI().toString(),
+                            managedIdentityRequest.headers,
+                            managedIdentityRequest.getBodyAsString());
             response = HttpHelper.executeHttpRequest(httpRequest, managedIdentityRequest.requestContext(), serviceBundle);
         } catch (URISyntaxException e) {
             throw new RuntimeException(e);
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/CloudShellManagedIdentitySource.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/CloudShellManagedIdentitySource.java
@@ -0,0 +1,88 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+import com.nimbusds.oauth2.sdk.util.URLUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.HashMap;
+
+class CloudShellManagedIdentitySource extends AbstractManagedIdentitySource{
+
+    private static final Logger LOG = LoggerFactory.getLogger(CloudShellManagedIdentitySource.class);
+
+    // MSI Constants. Docs for MSI are available here https://learn.microsoft.com/en-us/azure/cloud-shell/msi-authorization
+    private static URI endpointUri;
+
+    private URI endpoint;
+
+    @Override
+    public void createManagedIdentityRequest(String resource) {
+        managedIdentityRequest.baseEndpoint = endpoint;
+        managedIdentityRequest.method = HttpMethod.POST;
+
+        managedIdentityRequest.headers = new HashMap<>();
+        managedIdentityRequest.headers.put("ContentType", "application/x-www-form-urlencoded");
+        managedIdentityRequest.headers.put("Metadata", "true");
+
+        managedIdentityRequest.bodyParameters = new HashMap<>();
+        managedIdentityRequest.bodyParameters.put("resource", Collections.singletonList(resource));
+    }
+
+    private CloudShellManagedIdentitySource(MsalRequest msalRequest, ServiceBundle serviceBundle, URI endpoint)
+    {
+        super(msalRequest, serviceBundle, ManagedIdentitySourceType.CloudShell);
+        this.endpoint = endpoint;
+
+        ManagedIdentityIdType idType =
+                ((ManagedIdentityApplication) msalRequest.application()).getManagedIdentityId().getIdType();
+        if (idType != ManagedIdentityIdType.SystemAssigned) {
+            throw new MsalManagedIdentityException(MsalError.USER_ASSIGNED_MANAGED_IDENTITY_NOT_SUPPORTED,
+                    String.format(MsalErrorMessage.MANAGED_IDENTITY_USER_ASSIGNED_NOT_SUPPORTED, "cloud shell"),
+                    ManagedIdentitySourceType.CloudShell);
+        }
+    }
+
+    protected static AbstractManagedIdentitySource create(MsalRequest msalRequest, ServiceBundle serviceBundle) {
+
+        IEnvironmentVariables environmentVariables = getEnvironmentVariables((ManagedIdentityParameters) msalRequest.requestContext().apiParameters());
+        String msiEndpoint = environmentVariables.getEnvironmentVariable(Constants.MSI_ENDPOINT);
+
+
+        // if ONLY the env var MSI_ENDPOINT is set the MsiType is CloudShell
+        if (StringHelper.isNullOrBlank(msiEndpoint))
+        {
+            LOG.info("[Managed Identity] Cloud shell managed identity is unavailable.");
+            return null;
+        }
+
+        return validateEnvironmentVariables(msiEndpoint)
+                ? new CloudShellManagedIdentitySource(msalRequest, serviceBundle, endpointUri)
+                : null;
+    }
+
+    private static boolean validateEnvironmentVariables(String msiEndpoint)
+    {
+        endpointUri = null;
+
+        try
+        {
+            endpointUri = new URI(msiEndpoint);
+        }
+        catch (URISyntaxException ex)
+        {
+            throw new MsalManagedIdentityException(MsalError.INVALID_MANAGED_IDENTITY_ENDPOINT, String.format(
+                    MsalErrorMessage.MANAGED_IDENTITY_ENDPOINT_INVALID_URI_ERROR, "MSI_ENDPOINT", msiEndpoint, "Cloud Shell"),
+                    ManagedIdentitySourceType.CloudShell);
+        }
+
+        LOG.info("[Managed Identity] Environment variables validation passed for cloud shell managed identity. Endpoint URI: " + endpointUri + ". Creating cloud shell managed identity.");
+        return true;
+    }
+
+}
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IMDSManagedIdentitySource.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IMDSManagedIdentitySource.java
@@ -11,7 +11,6 @@
 import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Map;
 
 class IMDSManagedIdentitySource extends AbstractManagedIdentitySource{
 
@@ -125,7 +124,7 @@ public ManagedIdentityResponse handleResponse(
 
             message = message + " " + errorContentMessage;
 
-            LOG.error("Error message: {message} Http status code: {response.StatusCode}");
+            LOG.error(String.format("Error message: %s Http status code: %s"), message, response.statusCode());
             throw new MsalManagedIdentityException("managed_identity_request_failed", message,
                     ManagedIdentitySourceType.Imds);
         }
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityClient.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityClient.java
@@ -32,10 +32,12 @@ public ManagedIdentityResponse getManagedIdentityResponse(ManagedIdentityParamet
 
     // This method tries to create managed identity source for different sources, if none is created then defaults to IMDS.
     private static AbstractManagedIdentitySource createManagedIdentitySource(MsalRequest msalRequest,
-            ServiceBundle serviceBundle) throws Exception {
+            ServiceBundle serviceBundle) {
         AbstractManagedIdentitySource managedIdentitySource;
         if ((managedIdentitySource = AppServiceManagedIdentitySource.create(msalRequest, serviceBundle)) != null) {
             return managedIdentitySource;
+        } else if ((managedIdentitySource = CloudShellManagedIdentitySource.create(msalRequest, serviceBundle)) != null) {
+            return managedIdentitySource;
         } else {
             return new IMDSManagedIdentitySource(msalRequest, serviceBundle);
         }
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityRequest.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ManagedIdentityRequest.java
@@ -22,14 +22,21 @@ class ManagedIdentityRequest extends MsalRequest {
 
     Map<String, String> headers;
 
-    Map<String, String> bodyParameters;
+    Map<String, List<String>> bodyParameters;
 
     Map<String, List<String>> queryParameters;
 
     public ManagedIdentityRequest(ManagedIdentityApplication managedIdentityApplication, RequestContext requestContext) {
         super(managedIdentityApplication, requestContext);
     }
 
+    public String getBodyAsString() {
+        if (bodyParameters == null || bodyParameters.isEmpty())
+            return "";
+
+        return URLUtils.serializeParameters(bodyParameters);
+    }
+
     public URL computeURI() throws URISyntaxException {
         String endpoint = this.appendQueryParametersToBaseEndpoint();
         try {
@@ -40,7 +47,7 @@ public URL computeURI() throws URISyntaxException {
     }
 
     private String appendQueryParametersToBaseEndpoint() {
-        if (queryParameters.isEmpty()) {
+        if (queryParameters == null || queryParameters.isEmpty()) {
             return baseEndpoint.toString();
         }
 
diff --git a/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/MsalErrorMessage.java b/msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/MsalErrorMessage.java
@@ -15,7 +15,7 @@ class MsalErrorMessage {
 
     public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_CONFIGURABLE_AT_RUNTIME = "[Managed Identity] Service Fabric user assigned managed identity ClientId or ResourceId is not configurable at runtime.";
 
-    public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_SUPPORTED = "[Managed Identity] User assigned identity is not supported by the %s Managed Identity. To authenticate with the system assigned identity omit the client id in ManagedIdentityApplicationBuilder.Create().";
+    public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_SUPPORTED = "[Managed Identity] User assigned identity is not supported by the %s Managed Identity. To authenticate with the system assigned identity use ManagedIdentityApplication.builder(ManagedIdentityId.systemAssigned()).build().";
 
     public static final String SCOPES_REQUIRED = "At least one scope needs to be requested for this authentication flow. ";
 }
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTestDataProvider.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTestDataProvider.java
@@ -17,6 +17,10 @@ public static Stream<Arguments> createData() {
                         ManagedIdentityTests.resource),
                 Arguments.of(ManagedIdentitySourceType.AppService, ManagedIdentityTests.appServiceEndpoint,
                         ManagedIdentityTests.resourceDefaultSuffix),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        ManagedIdentityTests.resource),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        ManagedIdentityTests.resourceDefaultSuffix),
                 Arguments.of(ManagedIdentitySourceType.Imds, ManagedIdentityTests.IMDS_ENDPOINT,
                         ManagedIdentityTests.resource),
                 Arguments.of(ManagedIdentitySourceType.Imds, ManagedIdentityTests.IMDS_ENDPOINT,
@@ -37,12 +41,24 @@ public static Stream<Arguments> createDataUserAssigned() {
                         ManagedIdentityId.userAssignedResourceId(RESOURCE_ID)));
     }
 
+    public static Stream<Arguments> createDataUserAssignedNotSupported() {
+        return Stream.of(
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        ManagedIdentityId.userAssignedClientId(CLIENT_ID)),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        ManagedIdentityId.userAssignedResourceId(RESOURCE_ID)));
+    }
+
     public static Stream<Arguments> createDataWrongScope() {
         return Stream.of(
                 Arguments.of(ManagedIdentitySourceType.AppService, ManagedIdentityTests.appServiceEndpoint,
                         "user.read"),
                 Arguments.of(ManagedIdentitySourceType.AppService, ManagedIdentityTests.appServiceEndpoint,
                         "https://management.core.windows.net//user_impersonation"),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        "user.read"),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint,
+                        "https://management.core.windows.net//user_impersonation"),
                 Arguments.of(ManagedIdentitySourceType.Imds, ManagedIdentityTests.IMDS_ENDPOINT,
                         "user.read"),
                 Arguments.of(ManagedIdentitySourceType.Imds, ManagedIdentityTests.IMDS_ENDPOINT,
@@ -52,6 +68,7 @@ public static Stream<Arguments> createDataWrongScope() {
     public static Stream<Arguments> createDataError() {
         return Stream.of(
                 Arguments.of(ManagedIdentitySourceType.AppService, ManagedIdentityTests.appServiceEndpoint),
+                Arguments.of(ManagedIdentitySourceType.CloudShell, ManagedIdentityTests.cloudShellEndpoint),
                 Arguments.of(ManagedIdentitySourceType.Imds, ManagedIdentityTests.IMDS_ENDPOINT));
     }
 }
diff --git a/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java b/msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/ManagedIdentityTests.java
@@ -11,7 +11,6 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 
 import java.net.SocketException;
-import java.net.URISyntaxException;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Collections;
@@ -36,17 +35,15 @@ public class ManagedIdentityTests {
 
     private String getSuccessfulResponse(String resource) {
         long expiresOn = Instant.now().plus(1, ChronoUnit.HOURS).getEpochSecond();
-        String response = "{\"access_token\":\"accesstoken\",\"expires_on\":\"" + expiresOn + "\",\"resource\":\"" + resource + "\",\"token_type\":" +
+        return "{\"access_token\":\"accesstoken\",\"expires_on\":\"" + expiresOn + "\",\"resource\":\"" + resource + "\",\"token_type\":" +
                 "\"Bearer\",\"client_id\":\"client_id\"}";
-
-        return response;
     }
 
     private String getMsiErrorResponse() {
         return "{\"statusCode\":\"500\",\"message\":\"An unexpected error occured while fetching the AAD Token.\",\"correlationId\":\"7d0c9763-ff1d-4842-a3f3-6d49e64f4513\"}";
     }
 
-    private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource) throws URISyntaxException {
+    private HttpRequest expectedRequest(ManagedIdentitySourceType source, String resource) {
         return expectedRequest(source, resource, ManagedIdentityId.systemAssigned());
     }
 
@@ -55,17 +52,27 @@ private HttpRequest expectedRequest(ManagedIdentitySourceType source, String res
         String endpoint = null;
         Map<String, String> headers = new HashMap<>();
         Map<String, List<String>> queryParameters = new HashMap<>();
+        Map<String, List<String>> bodyParameters = new HashMap<>();
 
         switch (source) {
             case AppService: {
                 endpoint = appServiceEndpoint;
-                queryParameters = new HashMap<>();
+
                 queryParameters.put("api-version", Collections.singletonList("2019-08-01"));
                 queryParameters.put("resource", Collections.singletonList(resource));
-                headers = new HashMap<>();
+
                 headers.put("X-IDENTITY-HEADER", "secret");
                 break;
             }
+            case CloudShell: {
+                endpoint = cloudShellEndpoint;
+
+                headers.put("ContentType", "application/x-www-form-urlencoded");
+                headers.put("Metadata", "true");
+
+                bodyParameters.put("resource", Collections.singletonList(resource));
+                return new HttpRequest(HttpMethod.POST, computeUri(endpoint, queryParameters), headers, URLUtils.serializeParameters(bodyParameters));
+            }
             case Imds: {
                 endpoint = IMDS_ENDPOINT;
                 queryParameters.put("api-version", Collections.singletonList("2018-02-01"));
@@ -89,12 +96,12 @@ private HttpRequest expectedRequest(ManagedIdentitySourceType source, String res
 
     private String computeUri(String endpoint, Map<String, List<String>> queryParameters) {
         if (queryParameters.isEmpty()) {
-            return endpoint.toString();
+            return endpoint;
         }
 
         String queryString = URLUtils.serializeParameters(queryParameters);
 
-        return endpoint.toString() + "?" + queryString;
+        return endpoint + "?" + queryString;
     }
 
     private HttpResponse expectedResponse(int statusCode, String response) {
@@ -157,6 +164,35 @@ void managedIdentityTest_UserAssigned_SuccessfulResponse(ManagedIdentitySourceTy
         assertNotNull(result.accessToken());
     }
 
+    @ParameterizedTest
+    @MethodSource("com.microsoft.aad.msal4j.ManagedIdentityTestDataProvider#createDataUserAssignedNotSupported")
+    void managedIdentityTest_UserAssigned_NotSupported(ManagedIdentitySourceType source, String endpoint, ManagedIdentityId id) throws Exception {
+        IEnvironmentVariables environmentVariables = new EnvironmentVariablesHelper(source, endpoint);
+        DefaultHttpClient httpClientMock = mock(DefaultHttpClient.class);
+
+        ManagedIdentityApplication miApp = ManagedIdentityApplication
+                .builder(id)
+                .httpClient(httpClientMock)
+                .build();
+
+        try {
+            IAuthenticationResult result = miApp.acquireTokenForManagedIdentity(
+                    ManagedIdentityParameters.builder(resource)
+                            .environmentVariables(environmentVariables)
+                            .build()).get();
+        } catch (Exception e) {
+            assertNotNull(e);
+            assertInstanceOf(MsalManagedIdentityException.class, e.getCause());
+
+            MsalManagedIdentityException msalMsiException = (MsalManagedIdentityException) e.getCause();
+            assertEquals(ManagedIdentitySourceType.CloudShell, msalMsiException.managedIdentitySourceType);
+            assertEquals(MsalError.USER_ASSIGNED_MANAGED_IDENTITY_NOT_SUPPORTED, msalMsiException.errorCode());
+            return;
+        }
+
+        fail("MsalManagedIdentityException is expected but not thrown.");
+    }
+
     @ParameterizedTest
     @MethodSource("com.microsoft.aad.msal4j.ManagedIdentityTestDataProvider#createData")
     void managedIdentityTest_DifferentScopes_RequestsNewToken(ManagedIdentitySourceType source, String endpoint) throws Exception {

Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,7 @@ class MsalErrorMessage {`
`15`	`15`
`16`	`16`	`public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_CONFIGURABLE_AT_RUNTIME = "[Managed Identity] Service Fabric user assigned managed identity ClientId or ResourceId is not configurable at runtime.";`
`17`	`17`
`18`		`- public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_SUPPORTED = "[Managed Identity] User assigned identity is not supported by the %s Managed Identity. To authenticate with the system assigned identity omit the client id in ManagedIdentityApplicationBuilder.Create().";`
	`18`	`+ public static final String MANAGED_IDENTITY_USER_ASSIGNED_NOT_SUPPORTED = "[Managed Identity] User assigned identity is not supported by the %s Managed Identity. To authenticate with the system assigned identity use ManagedIdentityApplication.builder(ManagedIdentityId.systemAssigned()).build().";`
`19`	`19`
`20`	`20`	`public static final String SCOPES_REQUIRED = "At least one scope needs to be requested for this authentication flow. ";`
`21`	`21`	`}`