feat: allow arbitrary additional claims for JwtClaims

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

@@ -32,7 +32,9 @@
 package com.google.auth.oauth2;
 
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
 import java.io.Serializable;
+import java.util.Map;
 import javax.annotation.Nullable;
 
 /**
@@ -61,8 +63,15 @@ public abstract class JwtClaims implements Serializable {
   @Nullable
   abstract String getSubject();
 
+  /**
+   * Returns additional claims for this object. The returned map is not guaranteed to be mutable.
+   *
+   * @return additional claims
+   */
+  abstract Map<String, String> getAdditionalClaims();
+
   static Builder newBuilder() {
-    return new AutoValue_JwtClaims.Builder();
+    return new AutoValue_JwtClaims.Builder().setAdditionalClaims(ImmutableMap.<String, String>of());
   }
 
   /**
@@ -74,10 +83,15 @@ static Builder newBuilder() {
    * @return new claims
    */
   public JwtClaims merge(JwtClaims other) {
+    ImmutableMap.Builder<String, String> 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())
         .setSubject(other.getSubject() == null ? getSubject() : other.getSubject())
+        .setAdditionalClaims(newClaimsBuilder.build())
         .build();
   }
 
@@ -103,6 +117,8 @@ abstract static class Builder {
 
     abstract Builder setSubject(String subject);
 
+    abstract Builder setAdditionalClaims(Map<String, String> additionalClaims);
+
     abstract JwtClaims build();
   }
 }

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

@@ -114,6 +114,9 @@ public void refresh() throws IOException {
     payload.setIssuedAtTimeSeconds(currentTime / 1000);
     payload.setExpirationTimeSeconds(currentTime / 1000 + lifeSpanSeconds);
 
+    // Add all additional claims
+    payload.putAll(jwtClaims.getAdditionalClaims());
+
     synchronized (lock) {
       this.expiryInSeconds = payload.getExpirationTimeSeconds();
 

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

@@ -0,0 +1,136 @@
+/*
+ * Copyright 2019, Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *    * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *    * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *
+ *    * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package com.google.auth.oauth2;
+
+import static org.junit.Assert.*;
+
+import java.util.Collections;
+import java.util.Map;
+import org.junit.Test;
+
+public class JwtClaimsTest {
+
+  @Test
+  public void testMergeOverwritesFields() {
+    JwtClaims claims1 =
+        JwtClaims.newBuilder()
+            .setAudience("audience-1")
+            .setIssuer("issuer-1")
+            .setSubject("subject-1")
+            .build();
+    JwtClaims claims2 =
+        JwtClaims.newBuilder()
+            .setAudience("audience-2")
+            .setIssuer("issuer-2")
+            .setSubject("subject-2")
+            .build();
+    JwtClaims merged = claims1.merge(claims2);
+
+    assertEquals("audience-2", merged.getAudience());
+    assertEquals("issuer-2", merged.getIssuer());
+    assertEquals("subject-2", merged.getSubject());
+  }
+
+  @Test
+  public void testMergeDefaultValues() {
+    JwtClaims claims1 =
+        JwtClaims.newBuilder()
+            .setAudience("audience-1")
+            .setIssuer("issuer-1")
+            .setSubject("subject-1")
+            .build();
+    JwtClaims claims2 = JwtClaims.newBuilder().setAudience("audience-2").build();
+    JwtClaims merged = claims1.merge(claims2);
+
+    assertEquals("audience-2", merged.getAudience());
+    assertEquals("issuer-1", merged.getIssuer());
+    assertEquals("subject-1", merged.getSubject());
+  }
+
+  @Test
+  public void testMergeNull() {
+    JwtClaims claims1 = JwtClaims.newBuilder().build();
+    JwtClaims claims2 = JwtClaims.newBuilder().build();
+    JwtClaims merged = claims1.merge(claims2);
+
+    assertNull(merged.getAudience());
+    assertNull(merged.getIssuer());
+    assertNull(merged.getSubject());
+    assertNotNull(merged.getAdditionalClaims());
+    assertTrue(merged.getAdditionalClaims().isEmpty());
+  }
+
+  @Test
+  public void testEquals() {
+    JwtClaims claims1 =
+        JwtClaims.newBuilder()
+            .setAudience("audience-1")
+            .setIssuer("issuer-1")
+            .setSubject("subject-1")
+            .build();
+    JwtClaims claims2 =
+        JwtClaims.newBuilder()
+            .setAudience("audience-1")
+            .setIssuer("issuer-1")
+            .setSubject("subject-1")
+            .build();
+
+    assertEquals(claims1, claims2);
+  }
+
+  @Test
+  public void testAdditionalClaimsDefaults() {
+    JwtClaims claims = JwtClaims.newBuilder().build();
+    assertNotNull(claims.getAdditionalClaims());
+    assertTrue(claims.getAdditionalClaims().isEmpty());
+  }
+
+  @Test
+  public void testMergeAdditionalClaims() {
+    JwtClaims claims1 =
+        JwtClaims.newBuilder().setAdditionalClaims(Collections.singletonMap("foo", "bar")).build();
+    JwtClaims claims2 =
+        JwtClaims.newBuilder()
+            .setAdditionalClaims(Collections.singletonMap("asdf", "qwer"))
+            .build();
+    JwtClaims merged = claims1.merge(claims2);
+
+    assertNull(merged.getAudience());
+    assertNull(merged.getIssuer());
+    assertNull(merged.getSubject());
+    Map<String, String> mergedAdditionalClaims = merged.getAdditionalClaims();
+    assertNotNull(mergedAdditionalClaims);
+    assertEquals(2, mergedAdditionalClaims.size());
+    assertEquals("bar", mergedAdditionalClaims.get("foo"));
+    assertEquals("qwer", mergedAdditionalClaims.get("asdf"));
+  }
+}

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

@@ -44,6 +44,7 @@
 import com.google.auth.http.AuthHttpConstants;
 import java.io.IOException;
 import java.security.PrivateKey;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import org.junit.Test;
@@ -157,72 +158,6 @@ public void builder_requiresCompleteClaims() {
     }
   }
 
-  @Test
-  public void claims_merge_overwritesFields() {
-    JwtClaims claims1 =
-        JwtClaims.newBuilder()
-            .setAudience("audience-1")
-            .setIssuer("issuer-1")
-            .setSubject("subject-1")
-            .build();
-    JwtClaims claims2 =
-        JwtClaims.newBuilder()
-            .setAudience("audience-2")
-            .setIssuer("issuer-2")
-            .setSubject("subject-2")
-            .build();
-    JwtClaims merged = claims1.merge(claims2);
-
-    assertEquals("audience-2", merged.getAudience());
-    assertEquals("issuer-2", merged.getIssuer());
-    assertEquals("subject-2", merged.getSubject());
-  }
-
-  @Test
-  public void claims_merge_defaultValues() {
-    JwtClaims claims1 =
-        JwtClaims.newBuilder()
-            .setAudience("audience-1")
-            .setIssuer("issuer-1")
-            .setSubject("subject-1")
-            .build();
-    JwtClaims claims2 = JwtClaims.newBuilder().setAudience("audience-2").build();
-    JwtClaims merged = claims1.merge(claims2);
-
-    assertEquals("audience-2", merged.getAudience());
-    assertEquals("issuer-1", merged.getIssuer());
-    assertEquals("subject-1", merged.getSubject());
-  }
-
-  @Test
-  public void claims_merge_null() {
-    JwtClaims claims1 = JwtClaims.newBuilder().build();
-    JwtClaims claims2 = JwtClaims.newBuilder().build();
-    JwtClaims merged = claims1.merge(claims2);
-
-    assertNull(merged.getAudience());
-    assertNull(merged.getIssuer());
-    assertNull(merged.getSubject());
-  }
-
-  @Test
-  public void claims_equals() {
-    JwtClaims claims1 =
-        JwtClaims.newBuilder()
-            .setAudience("audience-1")
-            .setIssuer("issuer-1")
-            .setSubject("subject-1")
-            .build();
-    JwtClaims claims2 =
-        JwtClaims.newBuilder()
-            .setAudience("audience-1")
-            .setIssuer("issuer-1")
-            .setSubject("subject-1")
-            .build();
-
-    assertEquals(claims1, claims2);
-  }
-
   @Test
   public void jwtWithClaims_overwritesClaims() throws IOException {
     JwtClaims claims =
@@ -287,13 +222,56 @@ public void getRequestMetadata_hasJwtAccess() throws IOException {
     verifyJwtAccess(metadata, "some-audience", "some-issuer", "some-subject", PRIVATE_KEY_ID);
   }
 
+  @Test
+  public void getRequestMetadata_withAdditionalClaims_hasJwtAccess() throws IOException {
+    JwtClaims claims =
+        JwtClaims.newBuilder()
+            .setAudience("some-audience")
+            .setIssuer("some-issuer")
+            .setSubject("some-subject")
+            .setAdditionalClaims(Collections.singletonMap("foo", "bar"))
+            .build();
+    JwtCredentials credentials =
+        JwtCredentials.newBuilder()
+            .setJwtClaims(claims)
+            .setPrivateKey(getPrivateKey())
+            .setPrivateKeyId(PRIVATE_KEY_ID)
+            .build();
+
+    Map<String, List<String>> metadata = credentials.getRequestMetadata();
+    verifyJwtAccess(
+        metadata,
+        "some-audience",
+        "some-issuer",
+        "some-subject",
+        PRIVATE_KEY_ID,
+        Collections.singletonMap("foo", "bar"));
+  }
+
   private void verifyJwtAccess(
       Map<String, List<String>> metadata,
       String expectedAudience,
       String expectedIssuer,
       String expectedSubject,
       String expectedKeyId)
       throws IOException {
+    verifyJwtAccess(
+        metadata,
+        expectedAudience,
+        expectedIssuer,
+        expectedSubject,
+        expectedKeyId,
+        Collections.<String, String>emptyMap());
+  }
+
+  private void verifyJwtAccess(
+      Map<String, List<String>> metadata,
+      String expectedAudience,
+      String expectedIssuer,
+      String expectedSubject,
+      String expectedKeyId,
+      Map<String, String> expectedAdditionalClaims)
+      throws IOException {
     assertNotNull(metadata);
     List<String> authorizations = metadata.get(AuthHttpConstants.AUTHORIZATION);
     assertNotNull("Authorization headers not found", authorizations);
@@ -310,5 +288,9 @@ private void verifyJwtAccess(
     assertEquals(expectedSubject, signature.getPayload().getSubject());
     assertEquals(expectedAudience, signature.getPayload().getAudience());
     assertEquals(expectedKeyId, signature.getHeader().getKeyId());
+
+    for (Map.Entry<String, String> entry : expectedAdditionalClaims.entrySet()) {
+      assertEquals(entry.getValue(), signature.getPayload().get(entry.getKey()));
+    }
   }
 }