Extend type support on JWT claims

oauth2_http/java/com/google/auth/oauth2/JwtClaims.java

@@ -31,9 +31,14 @@
 
 package com.google.auth.oauth2;
 
+import com.google.api.client.util.Preconditions;
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.io.Serializable;
+import java.util.Collection;
+import java.util.Date;
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
 
@@ -68,10 +73,10 @@ public abstract class JwtClaims implements Serializable {
    *
    * @return additional claims
    */
-  abstract Map<String, String> getAdditionalClaims();
+  abstract Map<String, ?> getAdditionalClaims();
 
   public static Builder newBuilder() {
-    return new AutoValue_JwtClaims.Builder().setAdditionalClaims(ImmutableMap.<String, String>of());
+    return new AutoValue_JwtClaims.Builder().setAdditionalClaims(ImmutableMap.<String, Object>of());
   }
 
   /**
@@ -83,7 +88,7 @@ public static Builder newBuilder() {
    * @return new claims
    */
   public JwtClaims merge(JwtClaims other) {
-    ImmutableMap.Builder<String, String> newClaimsBuilder = ImmutableMap.builder();
+    ImmutableMap.Builder<String, Object> newClaimsBuilder = ImmutableMap.builder();
     newClaimsBuilder.putAll(getAdditionalClaims());
     newClaimsBuilder.putAll(other.getAdditionalClaims());
 
@@ -111,14 +116,123 @@ public boolean isComplete() {
 
   @AutoValue.Builder
   public abstract static class Builder {
+    /** Basic types supported by JSON standard. */
+    private static List<Class<? extends Serializable>> SUPPORTED_BASIC_TYPES =
+        ImmutableList.of(
+            String.class,
+            Integer.class,
+            Double.class,
+            Float.class,
+            Boolean.class,
+            Date.class,
+            String[].class,
+            Integer[].class,
+            Double[].class,
+            Float[].class,
+            Boolean[].class,
+            Date[].class);
+
+    private static final String ERROR_MESSAGE =
+        "Invalid type on additional claims. Valid types are String, Integer, "
+            + "Double, Float, Boolean, Date, List and Map. Map keys must be Strings.";
+
     public abstract Builder setAudience(String audience);
 
     public abstract Builder setIssuer(String issuer);
 
     public abstract Builder setSubject(String subject);
 
-    public abstract Builder setAdditionalClaims(Map<String, String> additionalClaims);
-
-    public abstract JwtClaims build();
+    public abstract Builder setAdditionalClaims(Map<String, ?> additionalClaims);
+
+    protected abstract JwtClaims autoBuild();
+
+    public JwtClaims build() {
+      JwtClaims claims = autoBuild();
+      Preconditions.checkState(validateClaims(claims.getAdditionalClaims()), ERROR_MESSAGE);
+      return claims;
+    }
+
+    /**
+     * Validate if the objects on a Map are valid for a JWT claim.
+     *
+     * @param claims Map of claim objects to be validated
+     */
+    private static boolean validateClaims(Map<String, ?> claims) {
+      if (!validateKeys(claims)) {
+        return false;
+      }
+
+      for (Object claim : claims.values()) {
+        if (!validateObject(claim)) {
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    /**
+     * Validates if the object is a valid JSON supported type.
+     *
+     * @param object to be evaluated
+     */
+    private static final boolean validateObject(@Nullable Object object) {
+      // According to JSON spec, null is a valid value.
+      if (object == null) {
+        return true;
+      }
+
+      if (object instanceof List) {
+        return validateCollection((List) object);
+      } else if (object instanceof Map) {
+        return validateKeys((Map) object) && validateCollection(((Map) object).values());
+      }
+
+      return isSupportedValue(object);
+    }
+
+    /**
+     * Validates the keys on a given map. Keys must be Strings.
+     *
+     * @param map map to be evaluated
+     */
+    private static final boolean validateKeys(Map map) {
+      for (Object key : map.keySet()) {
+        if (!(key instanceof String)) {
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    /**
+     * Validates if a collection is a valid JSON value. Empty collections are considered valid.
+     *
+     * @param collection collection to be evaluated
+     */
+    private static final boolean validateCollection(Collection collection) {
+      if (collection.isEmpty()) {
+        return true;
+      }
+
+      for (Object item : collection) {
+        if (!validateObject(item)) {
+          return false;
+        }
+      }
+
+      return true;
+    }
+
+    /**
+     * Validates if the given object is an instance of a valid JSON basic type.
+     *
+     * @param value object to be evaluated.
+     */
+    private static final boolean isSupportedValue(Object value) {
+      Class clazz = value.getClass();
+      return SUPPORTED_BASIC_TYPES.contains(clazz);
+    }
   }
 }

oauth2_http/javatests/com/google/auth/oauth2/JwtClaimsTest.java

@@ -31,11 +31,21 @@
 
 package com.google.auth.oauth2;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import java.util.Collections;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 
 public class JwtClaimsTest {
 
@@ -114,6 +124,43 @@ public void testAdditionalClaimsDefaults() {
     assertTrue(claims.getAdditionalClaims().isEmpty());
   }
 
+  @Test
+  public void testComplexAdditionalClaims() {
+    Map<String, String> map = ImmutableMap.of("aaa", "bbb");
+
+    Map<String, Object> complexClaims = new HashMap();
+    complexClaims.put("foo", "bar");
+    complexClaims.put("abc", 123);
+    complexClaims.put("def", 12.3);
+    complexClaims.put("ghi", map);
+    complexClaims.put("jkl", ImmutableList.of(1, 2, 3));
+    complexClaims.put("mno", new Date());
+
+    JwtClaims claims = JwtClaims.newBuilder().setAdditionalClaims(complexClaims).build();
+
+    Map<String, ?> additionalClaims = claims.getAdditionalClaims();
+    assertEquals(additionalClaims.size(), 6);
+    assertEquals(additionalClaims.get("ghi"), map);
+  }
+
+  @Test
+  public void testValidateAdditionalClaims() {
+    Map<String, Object> complexClaims = new HashMap<>();
+    complexClaims.put("abc", new HashSet<>());
+
+    final JwtClaims.Builder claimsBuilder =
+        JwtClaims.newBuilder().setAdditionalClaims(complexClaims);
+
+    assertThrows(
+        IllegalStateException.class,
+        new ThrowingRunnable() {
+          @Override
+          public void run() {
+            claimsBuilder.build();
+          }
+        });
+  }
+
   @Test
   public void testMergeAdditionalClaims() {
     JwtClaims claims1 =
@@ -124,13 +171,19 @@ public void testMergeAdditionalClaims() {
             .build();
     JwtClaims merged = claims1.merge(claims2);
 
+    Map<String, String> value = Collections.singletonMap("key2", "val2");
+    JwtClaims claims3 =
+        JwtClaims.newBuilder().setAdditionalClaims(Collections.singletonMap("def", value)).build();
+    JwtClaims complexMerged = merged.merge(claims3);
+
     assertNull(merged.getAudience());
     assertNull(merged.getIssuer());
     assertNull(merged.getSubject());
-    Map<String, String> mergedAdditionalClaims = merged.getAdditionalClaims();
+    Map<String, ?> mergedAdditionalClaims = complexMerged.getAdditionalClaims();
     assertNotNull(mergedAdditionalClaims);
-    assertEquals(2, mergedAdditionalClaims.size());
+    assertEquals(3, mergedAdditionalClaims.size());
     assertEquals("bar", mergedAdditionalClaims.get("foo"));
     assertEquals("qwer", mergedAdditionalClaims.get("asdf"));
+    assertEquals(value, mergedAdditionalClaims.get("def"));
   }
 }