Extend type support on JWT claims

Additional JWT claims now support more types other than String

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;
 
@@ -66,12 +71,13 @@ public abstract class JwtClaims implements Serializable {
   /**
    * Returns additional claims for this object. The returned map is not guaranteed to be mutable.
    *
-   * @return additional claims
+   * @return additional string 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,10 +89,11 @@ 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());
 
+
     return newBuilder()
         .setAudience(other.getAudience() == null ? getAudience() : other.getAudience())
         .setIssuer(other.getIssuer() == null ? getIssuer() : other.getIssuer())
@@ -111,14 +118,111 @@ 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);
+
+    public 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 c = value.getClass();
+      return SUPPORTED_BASIC_TYPES.contains(c);
+    }
   }
 }

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

@@ -34,6 +34,7 @@
 import com.google.api.client.json.webtoken.JsonWebSignature;
 import com.google.api.client.json.webtoken.JsonWebToken;
 import com.google.api.client.util.Clock;
+import com.google.api.client.util.GenericData;
 import com.google.auth.Credentials;
 import com.google.auth.http.AuthHttpConstants;
 import com.google.common.annotations.VisibleForTesting;

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

@@ -31,11 +31,15 @@
 
 package com.google.auth.oauth2;
 
+import static com.google.common.collect.ImmutableMap.builder;
 import static org.junit.Assert.*;
 
-import java.util.Collections;
-import java.util.Map;
+import java.util.*;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 
 public class JwtClaimsTest {
 
@@ -114,6 +118,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 +165,21 @@ 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"));
+    }
 }