Address PR feedback

Author: Avery-Dunn

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/ClientAssertion.java

@@ -12,7 +12,7 @@
 @EqualsAndHashCode
 final class ClientAssertion implements IClientAssertion {
 
-    static final String ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
+    static final String ASSERTION_TYPE_JWT_BEARER = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer";
     private final String assertion;
 
     ClientAssertion(final String assertion) {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/IBroker.java

@@ -4,8 +4,6 @@
 package com.microsoft.aad.msal4j;
 
 import java.net.URL;
-import java.nio.charset.StandardCharsets;
-import java.util.Base64;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -67,17 +65,10 @@ default IAuthenticationResult parseBrokerAuthResult(String authority, String idT
             if (idToken != null) {
                 builder.idToken(idToken);
                 if (accountId != null) {
-                    String idTokenJson;
-
-                    try {
-                        idTokenJson = new String(Base64.getDecoder().decode(idToken.split("\\.")[1]), StandardCharsets.UTF_8);
-                    } catch (ArrayIndexOutOfBoundsException e) {
-                        throw new MsalServiceException("Error parsing ID token, missing payload section. Ensure that the ID token is following the JWT format.",
-                                AuthenticationErrorCode.INVALID_JWT);
-                    }
+                    IdToken idTokenObj = JsonHelper.createIdTokenFromEncodedTokenString(idToken);
 
                     builder.accountCacheEntity(AccountCacheEntity.create(clientInfo,
-                            Authority.createAuthority(new URL(authority)), JsonHelper.convertJsonToObject(idTokenJson, IdToken.class), null));
+                            Authority.createAuthority(new URL(authority)), idTokenObj, null));
                 }
             }
             if (accessToken != null) {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/JsonHelper.java

@@ -15,10 +15,8 @@
 import com.fasterxml.jackson.databind.JsonNode;
 
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Set;
+import java.nio.charset.StandardCharsets;
+import java.util.*;
 
 class JsonHelper {
     static ObjectMapper mapper;
@@ -41,6 +39,18 @@ static <T> T convertJsonToObject(final String json, final Class<T> tClass) {
         }
     }
 
+    static IdToken createIdTokenFromEncodedTokenString(String token) {
+        String idTokenJson;
+        try {
+            idTokenJson = new String(Base64.getUrlDecoder().decode(token.split("\\.")[1]), StandardCharsets.UTF_8);
+        } catch (ArrayIndexOutOfBoundsException e) {
+            throw new MsalClientException("Error parsing ID token, missing payload section.",
+                    AuthenticationErrorCode.INVALID_JWT);
+        }
+
+        return JsonHelper.convertJsonToObject(idTokenJson, IdToken.class);
+    }
+
     //This method is used to convert a JSON string to an object which implements the JsonSerializable interface from com.azure.json
     static <T extends JsonSerializable<T>> T convertJsonStringToJsonSerializableObject(String jsonResponse, ReadValueCallback<JsonReader, T> readFunction) {
         try (JsonReader jsonReader = JsonProviders.createReader(jsonResponse)) {

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/JwtHelper.java

@@ -30,10 +30,7 @@ static ClientAssertion buildJwt(String clientId, final ClientCertificate credent
             header.put("typ", "JWT");
 
             if (sendX5c) {
-                List<String> certs = new ArrayList<>();
-                for (String cert : credential.getEncodedPublicKeyCertificateChain()) {
-                    certs.add(cert);
-                }
+                List<String> certs = new ArrayList<>(credential.getEncodedPublicKeyCertificateChain());
                 header.put("x5c", certs);
             }
 

msal4j-sdk/src/main/java/com/microsoft/aad/msal4j/TokenRequestExecutor.java

@@ -3,14 +3,11 @@
 
 package com.microsoft.aad.msal4j;
 
-import com.nimbusds.oauth2.sdk.ParseException;
-import com.nimbusds.oauth2.sdk.SerializeException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
 import java.net.MalformedURLException;
-import java.nio.charset.StandardCharsets;
 import java.util.*;
 
 class TokenRequestExecutor {
@@ -110,7 +107,7 @@ private void addQueryParameters(OAuthHttpRequest oauthHttpRequest) {
 
     private void addJWTBearerAssertionParams(Map<String, String> queryParameters, String assertion) {
         queryParameters.put("client_assertion", assertion);
-        queryParameters.put("client_assertion_type", "urn:ietf:params:oauth:client-assertion-type:jwt-bearer");
+        queryParameters.put("client_assertion_type", ClientAssertion.ASSERTION_TYPE_JWT_BEARER);
     }
 
     private AuthenticationResult createAuthenticationResultFromOauthHttpResponse(HttpResponse oauthHttpResponse) {
@@ -121,15 +118,7 @@ private AuthenticationResult createAuthenticationResultFromOauthHttpResponse(Htt
 
             AccountCacheEntity accountCacheEntity = null;
             if (!StringHelper.isNullOrBlank(response.idToken())) {
-                String idTokenJson;
-                try {
-                    idTokenJson = new String(Base64.getUrlDecoder().decode(response.idToken().split("\\.")[1]), StandardCharsets.UTF_8);
-                } catch (ArrayIndexOutOfBoundsException e) {
-                    throw new MsalServiceException("Error parsing ID token, missing payload section. Ensure that the ID token is following the JWT format.",
-                            AuthenticationErrorCode.INVALID_JWT);
-                }
-
-                IdToken idToken = JsonHelper.convertJsonToObject(idTokenJson, IdToken.class);
+                IdToken idToken = JsonHelper.createIdTokenFromEncodedTokenString(response.idToken());
 
                 AuthorityType type = msalRequest.application().authenticationAuthority.authorityType;
                 if (!StringHelper.isBlank(response.getClientInfo())) {

msal4j-sdk/src/test/java/com/microsoft/aad/msal4j/HelperAndUtilityTests.java

@@ -5,6 +5,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import java.nio.charset.StandardCharsets;
 import java.security.NoSuchAlgorithmException;
 import java.security.cert.CertificateEncodingException;
 import java.util.*;
@@ -147,4 +148,43 @@ void JwtHelper_buildJwt_ValidSha1AndSha256Assertions() throws MsalClientExceptio
         // Verify the correct certificate hash method was called
         verify(clientCertificateMock).publicCertificateHash();
     }
+
+    @Test
+    void JsonHelper_createIdTokenFromEncodedTokenString_Base64URLCharacters() {
+        HashMap<String, String> tokenParameters = new HashMap<>();
+        tokenParameters.put("preferred_username", "~nameWith~specialChars");
+        String encodedIDToken = TestHelper.createIdToken(tokenParameters);
+
+        try {
+            //TestHelper.createIdToken() should use Base64URL encoding, so first we prove that the encoded token cannot be decoded with Base64 decoder
+            Base64.getDecoder().decode(encodedIDToken);
+
+            fail("IllegalArgumentException was expected but not thrown.");
+        } catch (IllegalArgumentException e) {
+            //Encoded token should have some "-" characters in it
+            assertTrue(e.getMessage().contains("Illegal base64 character 2e"));
+        }
+
+        // Act
+        IdToken idToken = JsonHelper.createIdTokenFromEncodedTokenString(encodedIDToken);
+
+        // Assert
+        assertNotNull(idToken);
+        assertEquals("~nameWith~specialChars", idToken.preferredUsername);
+    }
+
+    @Test
+    void JsonHelper_createIdTokenFromEncodedTokenString_InvalidJsonInToken() {
+        // Arrange
+        String invalidPayload = "{not-valid-json}";
+        String encodedPayload = Base64.getUrlEncoder().withoutPadding()
+                .encodeToString(invalidPayload.getBytes(StandardCharsets.UTF_8));
+        String invalidToken = "header." + encodedPayload + ".signature";
+
+        // Act & Assert
+        MsalJsonParsingException exception = assertThrows(MsalJsonParsingException.class,
+                () -> JsonHelper.createIdTokenFromEncodedTokenString(invalidToken));
+
+        assertEquals(AuthenticationErrorCode.INVALID_JSON, exception.errorCode());
+    }
 }