Tenant Profiles

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,6 @@ class Account implements IAccount {
     String environment;
 
     String username;
+
+    Map<String, ?> idTokenClaims;
 }

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
@@ -49,6 +50,8 @@ ClientInfo clientInfo() {
     @JsonProperty("authority_type")
     protected String authorityType;
 
+    protected Map<String, ?> idTokenClaims;
+
     String getKey() {
 
         List<String> keyParts = new ArrayList<>();
@@ -77,6 +80,7 @@ static AccountCacheEntity create(String clientInfoStr, Authority requestAuthorit
             account.localAccountId(localAccountId);
             account.username(idToken.preferredUsername);
             account.name(idToken.name);
+            account.idTokenClaims(idToken.tokenClaims());
         }
 
         return account;
@@ -101,6 +105,10 @@ static AccountCacheEntity create(String clientInfoStr, Authority requestAuthorit
     }
 
     IAccount toAccount(){
-        return new Account(homeAccountId, environment, username);
+        if (homeAccountId.contains(localAccountId)) {
+            return new Account(homeAccountId, environment, username, idTokenClaims);
+        } else {
+            return new TenantProfile(homeAccountId, environment, username, idTokenClaims, localAccountId);
+        }
     }
 }

src/main/java/com/microsoft/aad/msal4j/IAccount.java

@@ -3,6 +3,7 @@
 
 package com.microsoft.aad.msal4j;
 
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -26,4 +27,9 @@ public interface IAccount {
      * @return account username
      */
     String username();
+
+    /**
+     * @return claims in ID token
+     */
+    Map<String, ?> idTokenClaims();
 }

src/main/java/com/microsoft/aad/msal4j/IMultiTenantAccount.java

@@ -0,0 +1,8 @@
+package com.microsoft.aad.msal4j;
+
+import java.util.Map;
+
+public interface IMultiTenantAccount extends IAccount {
+
+    Map<String, ITenantProfile> getTenantProfiles();
+}

src/main/java/com/microsoft/aad/msal4j/ITenantProfile.java

@@ -0,0 +1,5 @@
+package com.microsoft.aad.msal4j;
+
+public interface ITenantProfile extends IAccount {
+
+}

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 {
 
@@ -59,6 +61,29 @@ class IdToken {
     @JsonProperty("unique_name")
     protected String uniqueName;
 
+    /**
+     * Used to attach all claims in an ID token to an account object
+     *
+     * @return map of key/value pairs of claims in ID token
+     */
+    protected Map<String, Object> tokenClaims() {
+        Map<String, Object> idTokenClaims = new HashMap<>();
+        idTokenClaims.put(ISSUER, this.issuer);
+        idTokenClaims.put(SUBJECT, this.subject);
+        idTokenClaims.put(AUDIENCE, this.audience);
+        idTokenClaims.put(EXPIRATION_TIME, this.expirationTime);
+        idTokenClaims.put(ISSUED_AT, this.issuedAt);
+        idTokenClaims.put(NOT_BEFORE, this.notBefore);
+        idTokenClaims.put(NAME, this.name);
+        idTokenClaims.put(PREFERRED_USERNAME, this.preferredUsername);
+        idTokenClaims.put(OBJECT_IDENTIFIER, this.objectIdentifier);
+        idTokenClaims.put(TENANT_IDENTIFIER, this.tenantIdentifier);
+        idTokenClaims.put(UPN, this.upn);
+        idTokenClaims.put(UNIQUE_NAME, this.uniqueName);
+
+        return idTokenClaims;
+    }
+
     static IdToken createFromJWTClaims(final JWTClaimsSet claims) throws ParseException {
         IdToken idToken = new IdToken();
 

src/main/java/com/microsoft/aad/msal4j/MultiTenantAccount.java

@@ -0,0 +1,22 @@
+package com.microsoft.aad.msal4j;
+
+import lombok.Getter;
+import lombok.Setter;
+
+import java.util.Map;
+
+@Getter
+@Setter
+public class MultiTenantAccount extends Account implements IMultiTenantAccount {
+
+    private Map<String, ITenantProfile> tenantProfiles;
+
+    public MultiTenantAccount(String homeAccountId, String environment, String username, Map<String, ?> idTokenClaims) {
+        super(homeAccountId, environment, username, idTokenClaims);
+    }
+
+    @Override
+    public Map<String, ITenantProfile> getTenantProfiles() {
+        return tenantProfiles;
+    }
+}

src/main/java/com/microsoft/aad/msal4j/TenantProfile.java

@@ -0,0 +1,20 @@
+package com.microsoft.aad.msal4j;
+
+import lombok.Getter;
+import lombok.Setter;
+import lombok.experimental.Accessors;
+
+import java.util.Map;
+
+@Accessors(fluent = true)
+@Getter
+@Setter
+public class TenantProfile extends Account implements ITenantProfile{
+
+    String tenantId;
+
+    public TenantProfile(String homeAccountId, String environment, String username, Map<String, ?> idTokenClaims, String tenantId) {
+        super(homeAccountId, environment, username, idTokenClaims);
+        this.tenantId = tenantId;
+    }
+}

src/main/java/com/microsoft/aad/msal4j/TokenCache.java

@@ -300,10 +300,56 @@ Set<IAccount> getAccounts(String clientId, Set<String> environmentAliases) {
                         build())) {
             try {
                 lock.readLock().lock();
+                Set<IAccount> rootAccounts = new HashSet<>();
 
-                return accounts.values().stream().
+                //Filter accounts map into two sets: a set of home accounts (local UID = home UID),
+                // and a set of tenant profiles (local UID != home UID)
+                Set<IAccount> homeAccounts = accounts.values().stream().
                         filter(acc -> environmentAliases.contains(acc.environment())).
+                        filter(acc -> acc.homeAccountId().contains(acc.localAccountId())).
                         collect(Collectors.mapping(AccountCacheEntity::toAccount, Collectors.toSet()));
+                Set<IAccount> tenantAccounts = accounts.values().stream().
+                        filter(acc -> environmentAliases.contains(acc.environment())).
+                        filter(acc -> !acc.homeAccountId().contains(acc.localAccountId())).
+                        collect(Collectors.mapping(AccountCacheEntity::toAccount, Collectors.toSet()));
+
+                Iterator<IAccount> homeAcctsIterator = homeAccounts.iterator();
+                Iterator<IAccount> tenantAcctsIterator;
+                Map<String, ITenantProfile> tenantProfiles;
+                IAccount homeAcc, tenantAcc;
+
+                while (homeAcctsIterator.hasNext()) {
+                    homeAcc = homeAcctsIterator.next();
+                    tenantAcctsIterator = tenantAccounts.iterator();
+                    tenantProfiles = new HashMap<>();
+
+                    while (tenantAcctsIterator.hasNext()) {
+                        tenantAcc = tenantAcctsIterator.next();
+
+                        //If the tenant account's home ID matches a home ID (UID) in the list of home accounts, it is a
+                        //tenant profile of that home account
+                        if (tenantAcc.homeAccountId().contains(homeAcc.homeAccountId().substring(0, homeAcc.homeAccountId().indexOf(".")))) {
+                            TenantProfile profile = (TenantProfile) tenantAcc;
+                            tenantProfiles.put(profile.tenantId(), profile);
+                            tenantAcctsIterator.remove();
+                        }
+                    }
+                    if (tenantProfiles.isEmpty()) {
+                        //If the home account has no tenant profiles, leave it as an Account object
+                        rootAccounts.add(homeAcc);
+                    } else {
+                        //If the home account had some tenant profiles, transform the home Account object into a
+                        //MultiTenantAccount object and attach its tenant profile list
+                        MultiTenantAccount multAcct = new MultiTenantAccount(
+                                homeAcc.homeAccountId(), homeAcc.environment(), homeAcc.username(), idTokens);
+                        multAcct.setTenantProfiles(tenantProfiles);
+                        rootAccounts.add(multAcct);
+                    }
+                }
+                //Add any tenant accounts which did not have a corresponding home account to the list of root accounts
+                rootAccounts.addAll(tenantAccounts);
+
+                return rootAccounts;
             } finally {
                 lock.readLock().unlock();
             }

src/test/java/com/microsoft/aad/msal4j/AccountTest.java

@@ -0,0 +1,76 @@
+// 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_ClassTypes() throws IOException, URISyntaxException {
+
+        ITokenCacheAccessAspect accountCache = new CachePersistenceIT.TokenPersistence(
+                TestHelper.readResource(this.getClass(),
+                        "/cache_data/multi-tenant-account-cache.json"));
+
+        PublicClientApplication app = PublicClientApplication.builder("uid1")
+                .setTokenCacheAccessAspect(accountCache).build();
+
+        Assert.assertEquals(app.getAccounts().join().size(), 3);
+        Iterator<IAccount> acctIterator = app.getAccounts().join().iterator();
+
+        IAccount curAccount;
+        IAccount account = null, multiAccount = null, tenantProfileAccount = null;
+        while (acctIterator.hasNext()) {
+            curAccount = acctIterator.next();
+            switch (curAccount.username()) {
+                case "Account":
+                    account = curAccount;
+                    break;
+                case "MultiAccount":
+                    multiAccount = curAccount;
+                    break;
+                case "TenantProfileNoHome":
+                    tenantProfileAccount = curAccount;
+                    break;
+            }
+        }
+
+        Assert.assertTrue(account instanceof Account);
+        Assert.assertTrue(multiAccount instanceof MultiTenantAccount);
+        Assert.assertTrue(tenantProfileAccount instanceof TenantProfile);
+    }
+
+    @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("uid1")
+                .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();
+
+            if (curAccount.username().equals("MultiAccount")) {
+                Assert.assertTrue(curAccount instanceof MultiTenantAccount);
+                Map<String, ITenantProfile> tenantProfiles = ((MultiTenantAccount) curAccount).getTenantProfiles();
+                Assert.assertNotNull(tenantProfiles);
+                Assert.assertNotNull(tenantProfiles.get("tenantprofile1"));
+                Assert.assertEquals(tenantProfiles.get("tenantprofile1").username(), "TenantProfile");
+            }
+        }
+    }
+}

src/test/resources/cache_data/multi-tenant-account-cache.json

@@ -0,0 +1,37 @@
+{
+  "Account": {
+    "uid1.utid1-login.windows.net-contoso": {
+      "username": "MultiAccount",
+      "local_account_id": "uid1",
+      "realm": "contoso",
+      "environment": "login.windows.net",
+      "home_account_id": "uid1.utid1",
+      "authority_type": "MSSTS"
+    },
+    "uid1.utid2-login.windows.net-contoso": {
+      "username": "TenantProfile",
+      "local_account_id": "tenantprofile1",
+      "realm": "contoso",
+      "environment": "login.windows.net",
+      "home_account_id": "uid1.utid1",
+      "authority_type": "MSSTS"
+    },
+    "uid2.utid2-login.windows.net-contoso": {
+      "username": "TenantProfileNoHome",
+      "local_account_id": "tenantprofile2",
+      "realm": "contoso",
+      "environment": "login.windows.net",
+      "home_account_id": "uid2.utid2",
+      "authority_type": "MSSTS"
+    },
+    "uid3.utid3-login.windows.net-contoso": {
+      "username": "Account",
+      "local_account_id": "uid3",
+      "realm": "contoso",
+      "environment": "login.windows.net",
+      "home_account_id": "uid3.utid3",
+      "authority_type": "MSSTS"
+    }
+  },
+  "AppMetadata": {}
+}