Tenant Profiles (#263)

* Classes for tenant profile functionality * Implement tenant profile feature * Tests for tenant profile feature * Simplify tenant profile class structure * 1.6.2 release * Classes for tenant profile redesign * Tests for tenant profile redesign * Adjust sample cached ID tokens to have realistic headers * Redesign how Tenant Pofiles are added to Accounts * New error code for JWT parse exceptions * Add claims and tenant profiles fields to Account * Remove annotation excluding realm field from comparisons * Use more generic token * Remove ID token claims field from Account * Minor changes for clarity * Adjust tests for tenant profile design refactor * Refactor tenant profile structure * Minor fixes * Minor fixes * Minor fixes * Simplify tenant profile class Co-authored-by: SomkaPe <pesomka@microsoft.com>
AzureAD · Aug 31, 2020 · 0b20b14 · 0b20b14
1 parent fce61b2
commit 0b20b14
Show file tree

Hide file tree

Showing 13 changed files with 280 additions and 8 deletions.
diff --git a/src/integrationtest/java/com.microsoft.aad.msal4j/TokenCacheIT.java b/src/integrationtest/java/com.microsoft.aad.msal4j/TokenCacheIT.java
@@ -146,7 +146,7 @@ public void twoAccountsInCache_SameUserDifferentTenants_RemoveAccountTest() thro
                 .get();
 
         // There should be two tokens in cache, with same accounts except for tenant
-        Assert.assertEquals(pca2.getAccounts().join().size() , 2);
+        Assert.assertEquals(pca2.getAccounts().join().iterator().next().getTenantProfiles().size() , 2);
 
         IAccount account = pca2.getAccounts().get().iterator().next();
 

diff --git a/src/main/java/com/microsoft/aad/msal4j/Account.java b/src/main/java/com/microsoft/aad/msal4j/Account.java
@@ -7,6 +7,7 @@
 import lombok.Getter;
 import lombok.Setter;
 import lombok.experimental.Accessors;
+import java.util.Map;
 
 /**
  * Representation of a single user account. If modifying this object, ensure it is compliant with
@@ -23,4 +24,10 @@ class Account implements IAccount {
     String environment;
 
     String username;
+
+    Map<String, ITenantProfile> tenantProfiles;
+
+    public Map<String, ITenantProfile> getTenantProfiles() {
+        return tenantProfiles;
+    }
 }
diff --git a/src/main/java/com/microsoft/aad/msal4j/AccountCacheEntity.java b/src/main/java/com/microsoft/aad/msal4j/AccountCacheEntity.java
@@ -10,6 +10,7 @@
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 
 @Accessors(fluent = true)
 @Getter
@@ -26,7 +27,6 @@ class AccountCacheEntity implements Serializable {
     @JsonProperty("environment")
     protected String environment;
 
-    @EqualsAndHashCode.Exclude
     @JsonProperty("realm")
     protected String realm;
 
@@ -101,6 +101,6 @@ static AccountCacheEntity create(String clientInfoStr, Authority requestAuthorit
     }
 
     IAccount toAccount(){
-        return new Account(homeAccountId, environment, username);
+        return new Account(homeAccountId, environment, username, null);
     }
 }
diff --git a/src/main/java/com/microsoft/aad/msal4j/AuthenticationErrorCode.java b/src/main/java/com/microsoft/aad/msal4j/AuthenticationErrorCode.java
@@ -100,4 +100,9 @@ public class AuthenticationErrorCode {
      * A JSON processing failure, indicating the JSON provided to MSAL is of invalid format.
      */
     public final static String INVALID_JSON = "invalid_json";
+
+    /**
+     * A JWT parsing failure, indicating the JWT provided to MSAL is of invalid format.
+     */
+    public final static String INVALID_JWT = "invalid_jwt";
 }
diff --git a/src/main/java/com/microsoft/aad/msal4j/IAccount.java b/src/main/java/com/microsoft/aad/msal4j/IAccount.java
@@ -3,7 +3,7 @@
 
 package com.microsoft.aad.msal4j;
 
-import java.util.Set;
+import java.util.Map;
 
 /**
  * Interface representing a single user account. An IAccount is returned in the {@link IAuthenticationResult}
@@ -26,4 +26,12 @@ public interface IAccount {
      * @return account username
      */
     String username();
+
+    /**
+     * Map of {@link ITenantProfile} objects related to this account, the keys of the map are the tenant ID values and
+     * match the 'realm' key of an ID token
+     *
+     * @return tenant profiles
+     */
+    Map<String, ITenantProfile> getTenantProfiles();
 }
diff --git a/src/main/java/com/microsoft/aad/msal4j/ITenantProfile.java b/src/main/java/com/microsoft/aad/msal4j/ITenantProfile.java
@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+import java.util.Map;
+
+/**
+ * Interface representing a single tenant profile. ITenantProfiles are made available through the
+ * {@link IAccount#getTenantProfiles()} method of an Account
+ *
+ */
+public interface ITenantProfile {
+
+    /**
+     * A map of claims taken from an ID token. Keys and values will follow the structure of a JSON Web Token
+     *
+     * @return Map claims in id token
+     */
+    Map<String, ?> getClaims();
+
+}
diff --git a/src/main/java/com/microsoft/aad/msal4j/IdToken.java b/src/main/java/com/microsoft/aad/msal4j/IdToken.java
@@ -7,6 +7,8 @@
 import com.nimbusds.jwt.JWTClaimsSet;
 
 import java.text.ParseException;
+import java.util.HashMap;
+import java.util.Map;
 
 class IdToken {
 

diff --git a/src/main/java/com/microsoft/aad/msal4j/TenantProfile.java b/src/main/java/com/microsoft/aad/msal4j/TenantProfile.java
@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.experimental.Accessors;
+import java.util.Map;
+
+/**
+ * Representation of a single tenant profile
+ */
+@Accessors(fluent = true)
+@Getter
+@Setter
+@AllArgsConstructor
+class TenantProfile implements ITenantProfile {
+
+    Map<String, ?> idTokenClaims;
+
+    public Map<String, ?> getClaims() {
+        return idTokenClaims;
+    }
+}
diff --git a/src/main/java/com/microsoft/aad/msal4j/TokenCache.java b/src/main/java/com/microsoft/aad/msal4j/TokenCache.java
@@ -6,8 +6,10 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.nimbusds.jwt.JWTParser;
 
 import java.io.IOException;
+import java.text.ParseException;
 import java.util.*;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -300,16 +302,63 @@ Set<IAccount> getAccounts(String clientId, Set<String> environmentAliases) {
                         build())) {
             try {
                 lock.readLock().lock();
+                Map<String, IAccount> rootAccounts = new HashMap<>();
 
-                return accounts.values().stream().
+                Set<AccountCacheEntity> accountsCached = accounts.values().stream().
                         filter(acc -> environmentAliases.contains(acc.environment())).
-                        collect(Collectors.mapping(AccountCacheEntity::toAccount, Collectors.toSet()));
+                        collect(Collectors.toSet());
+
+                for (AccountCacheEntity accCached : accountsCached) {
+
+                    IdTokenCacheEntity idToken = idTokens.get(getIdTokenKey(
+                            accCached.homeAccountId(),
+                            accCached.environment(),
+                            clientId,
+                            accCached.realm()));
+
+                    Map<String, ?> idTokenClaims = null;
+                    if (idToken != null) {
+                        idTokenClaims = JWTParser.parse(idToken.secret()).getJWTClaimsSet().getClaims();
+                    }
+
+                    ITenantProfile profile = new TenantProfile(idTokenClaims);
+
+                    if (rootAccounts.get(accCached.homeAccountId()) == null) {
+                        IAccount acc = accCached.toAccount();
+                        ((Account) acc).tenantProfiles = new HashMap<>();
+                        ((Account) acc).tenantProfiles().put(accCached.realm(), profile);
+
+                        rootAccounts.put(accCached.homeAccountId(), acc);
+                    } else {
+                        ((Account)rootAccounts.get(accCached.homeAccountId())).tenantProfiles.put(accCached.realm(), profile);
+                        if (accCached.homeAccountId().contains(accCached.localAccountId())) {
+                            ((Account) rootAccounts.get(accCached.homeAccountId())).username(accCached.username());
+                        }
+                    }
+                }
+
+                return new HashSet<>(rootAccounts.values());
+            } catch (ParseException e) {
+                throw new MsalClientException("Cached JWT could not be parsed: " + e.getMessage(), AuthenticationErrorCode.INVALID_JWT);
             } finally {
                 lock.readLock().unlock();
             }
         }
     }
 
+    /**
+     * Returns a String representing a key of a cached ID token, formatted in the same way as {@link IdTokenCacheEntity#getKey}
+     *
+     * @return String representing a possible key of a cached ID token
+     */
+    private String getIdTokenKey(String homeAccountId, String environment, String clientId, String realm) {
+        return String.join(Constants.CACHE_KEY_SEPARATOR,
+                Arrays.asList(homeAccountId,
+                        environment,
+                        "idtoken", clientId,
+                        realm, "")).toLowerCase();
+    }
+
     /**
      * @return familyId status of application
      */

diff --git a/src/samples/cache/sample_cache.json b/src/samples/cache/sample_cache.json
@@ -38,7 +38,7 @@
       "realm": "contoso",
       "environment": "login.windows.net",
       "credential_type": "IdToken",
-      "secret": "header.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
       "client_id": "my_client_id",
       "home_account_id": "uid.utid"
     }

diff --git a/src/test/java/com/microsoft/aad/msal4j/AccountTest.java b/src/test/java/com/microsoft/aad/msal4j/AccountTest.java
@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.microsoft.aad.msal4j;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.util.Iterator;
+import java.util.Map;
+
+public class AccountTest {
+
+    @Test
+    public void testMultiTenantAccount_AccessTenantProfile() throws IOException, URISyntaxException {
+
+        ITokenCacheAccessAspect accountCache = new CachePersistenceIT.TokenPersistence(
+                TestHelper.readResource(this.getClass(),
+                        "/cache_data/multi-tenant-account-cache.json"));
+
+        PublicClientApplication app = PublicClientApplication.builder("client_id")
+                .setTokenCacheAccessAspect(accountCache).build();
+
+        Assert.assertEquals(app.getAccounts().join().size(), 3);
+        Iterator<IAccount> acctIterator = app.getAccounts().join().iterator();
+
+        IAccount curAccount;
+        while (acctIterator.hasNext()) {
+            curAccount = acctIterator.next();
+
+            switch (curAccount.username()) {
+                case "MultiTenantAccount": {
+                    Assert.assertEquals(curAccount.homeAccountId(), "uid1.utid1");
+                    Map<String, ITenantProfile> tenantProfiles = curAccount.getTenantProfiles();
+                    Assert.assertNotNull(tenantProfiles);
+                    Assert.assertEquals(tenantProfiles.size(), 3);
+                    Assert.assertNotNull(tenantProfiles.get("utid1"));
+                    Assert.assertNotNull(tenantProfiles.get("utid1").getClaims());
+                    Assert.assertNotNull(tenantProfiles.get("utid2"));
+                    Assert.assertNotNull(tenantProfiles.get("utid2").getClaims());
+                    Assert.assertNotNull(tenantProfiles.get("utid3"));
+                    Assert.assertNotNull(tenantProfiles.get("utid3").getClaims());
+                    break;
+                }
+                case "SingleTenantAccount": {
+                    Assert.assertEquals(curAccount.homeAccountId(), "uid6.utid5");
+                    Map<String, ITenantProfile> tenantProfiles = curAccount.getTenantProfiles();
+                    Assert.assertNotNull(tenantProfiles);
+                    Assert.assertEquals(tenantProfiles.size(), 1);
+                    Assert.assertNotNull(tenantProfiles.get("utid5"));
+                    Assert.assertNotNull(tenantProfiles.get("utid5").getClaims());
+                    break;
+                }
+                case "TenantProfileNoHome": {
+                    Assert.assertEquals(curAccount.homeAccountId(), "uid5.utid4");
+                    Map<String, ITenantProfile> tenantProfiles = curAccount.getTenantProfiles();
+                    Assert.assertNotNull(tenantProfiles);
+                    Assert.assertEquals(tenantProfiles.size(), 1);
+                    Assert.assertNotNull(tenantProfiles.get("utid4"));
+                    Assert.assertNotNull(tenantProfiles.get("utid4").getClaims());
+                    break;
+                }
+            }
+        }
+    }
+}
diff --git a/src/test/resources/cache_data/multi-tenant-account-cache.json b/src/test/resources/cache_data/multi-tenant-account-cache.json
@@ -0,0 +1,86 @@
+{
+  "Account": {
+    "uid1.utid1-login.windows.net-utid1": {
+      "username": "MultiTenantAccount",
+      "local_account_id": "uid1",
+      "realm": "utid1",
+      "environment": "login.windows.net",
+      "home_account_id": "uid1.utid1",
+      "authority_type": "MSSTS"
+    },
+    "uid1.utid1-login.windows.net-utid2": {
+      "username": "TenantProfile1",
+      "local_account_id": "uid2",
+      "realm": "utid2",
+      "environment": "login.windows.net",
+      "home_account_id": "uid1.utid1",
+      "authority_type": "MSSTS"
+    },
+    "uid1.utid1-login.windows.net-utid3": {
+      "username": "TenantProfile2",
+      "local_account_id": "uid3",
+      "realm": "utid3",
+      "environment": "login.windows.net",
+      "home_account_id": "uid1.utid1",
+      "authority_type": "MSSTS"
+    },
+    "uid5.utid4-login.windows.net-utid4": {
+      "username": "TenantProfileNoHome",
+      "local_account_id": "uid4",
+      "realm": "utid4",
+      "environment": "login.windows.net",
+      "home_account_id": "uid5.utid4",
+      "authority_type": "MSSTS"
+    },
+    "uid6.utid5-login.windows.net-utid5": {
+      "username": "SingleTenantAccount",
+      "local_account_id": "uid6",
+      "realm": "utid5",
+      "environment": "login.windows.net",
+      "home_account_id": "uid6.utid5",
+      "authority_type": "MSSTS"
+    }
+  },
+  "IdToken": {
+    "uid1.utid1-login.windows.net-idtoken-client_id-utid1-": {
+      "realm": "utid1",
+      "environment": "login.windows.net",
+      "credential_type": "IdToken",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "client_id": "client_id",
+      "home_account_id": "uid.utid1"
+    },
+    "uid1.utid1-login.windows.net-idtoken-client_id-utid2-": {
+      "realm": "utid2",
+      "environment": "login.windows.net",
+      "credential_type": "IdToken",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "client_id": "client_id",
+      "home_account_id": "uid.utid1"
+    },
+    "uid1.utid1-login.windows.net-idtoken-client_id-utid3-": {
+      "realm": "utid3",
+      "environment": "login.windows.net",
+      "credential_type": "IdToken",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "client_id": "client_id",
+      "home_account_id": "uid.utid1"
+    },
+    "uid5.utid4-login.windows.net-idtoken-client_id-utid4-": {
+      "realm": "utid4",
+      "environment": "login.windows.net",
+      "credential_type": "IdToken",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "client_id": "client_id",
+      "home_account_id": "uid5.utid4"
+    },
+    "uid6.utid5-login.windows.net-idtoken-client_id-utid5-": {
+      "realm": "utid5",
+      "environment": "login.windows.net",
+      "credential_type": "IdToken",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "client_id": "client_id",
+      "home_account_id": "uid6.utid5"
+    }
+  }
+}
diff --git a/src/test/resources/cache_data/serialized_cache.json b/src/test/resources/cache_data/serialized_cache.json
@@ -41,7 +41,7 @@
       "realm": "contoso",
       "environment": "login.windows.net",
       "credential_type": "IdToken",
-      "secret": "header.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
+      "secret": "eyJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiIsImtpZCI6Imh1Tjk1SXZQZmVocTM0R3pCRFoxR1hHaXJuTSJ9.eyJvaWQiOiAib2JqZWN0MTIzNCIsICJwcmVmZXJyZWRfdXNlcm5hbWUiOiAiSm9obiBEb2UiLCAic3ViIjogInN1YiJ9.signature",
       "client_id": "my_client_id",
       "home_account_id": "uid.utid"
     }