Make UserPrincipalManager#getRoles more robust. (Azure#31803)

Author: rujche

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/aad/filter/UserPrincipalManager.java

@@ -30,12 +30,13 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.text.ParseException;
-import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 /**
  * A user principal manager to load user info from JWT.
@@ -153,11 +154,19 @@ public UserPrincipal buildUserPrincipal(String aadIssuedBearerToken) throws Pars
     }
 
     Set<String> getRoles(JWTClaimsSet set) {
-        return Optional.of(set)
-                .map(p -> p.getClaim(AadJwtClaimNames.ROLES))
-                .map(Collection.class::cast)
-                .map(Collection<Object>::stream)
-                .orElseGet(Stream::empty)
+        if (set == null) {
+            return Collections.emptySet();
+        }
+        Object rolesClaim = set.getClaim(AadJwtClaimNames.ROLES);
+        if (rolesClaim == null) {
+            return Collections.emptySet();
+        }
+        if (rolesClaim instanceof Iterable<?>) {
+            return StreamSupport.stream(((Iterable<?>) rolesClaim).spliterator(), false)
+                    .map(Object::toString)
+                    .collect(Collectors.toSet());
+        }
+        return Stream.of(rolesClaim)
                 .map(Object::toString)
                 .collect(Collectors.toSet());
     }

sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/aad/filter/UserPrincipalManagerTests.java

@@ -20,7 +20,10 @@
 import java.nio.file.Paths;
 import java.security.cert.CertificateFactory;
 import java.security.cert.X509Certificate;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.Set;
 import java.util.stream.Stream;
 
@@ -77,14 +80,21 @@ void nullIssuer() {
     }
 
     @Test
-    void testRolesExtracted() {
+    void getRolesTest() {
+        rolesExtractedAsExpected(null, new ArrayList<>());
+        rolesExtractedAsExpected("role1", Arrays.asList("role1"));
+        rolesExtractedAsExpected(Arrays.asList("role1", "role2"), Arrays.asList("role1", "role2"));
+        rolesExtractedAsExpected(new HashSet<>(Arrays.asList("role1", "role2")), Arrays.asList("role1", "role2"));
+    }
+
+    private void rolesExtractedAsExpected(Object rolesClaimValue, Collection<String> expected) {
         JWTClaimsSet set = new JWTClaimsSet.Builder()
-                .claim("roles", Arrays.asList("role1", "role2"))
+                .claim("roles", rolesClaimValue)
                 .build();
-        Set<String> result = new UserPrincipalManager(null).getRoles(set);
-        assertEquals(2, result.size());
-        assertTrue(result.contains("role1"));
-        assertTrue(result.contains("role2"));
+        Set<String> actual = new UserPrincipalManager(null).getRoles(set);
+        assertEquals(expected.size(), actual.size());
+        assertTrue(expected.containsAll(actual));
+        assertTrue(actual.containsAll(expected));
     }
 
     private String readJwtValidIssuerTxt() {