From d28a6ed08a5236b4aa3ff1e6249b9437f574f243 Mon Sep 17 00:00:00 2001
From: Jeff Ching <chingor@google.com>
Date: Thu, 22 Aug 2019 10:32:54 -0700
Subject: [PATCH] fix: allow unset/null privateKeyId for JwtCredentials (#336)

* fix: add failing test for unset/null privateKeyId

* fix: relax preconditions to allow null privateKeyId

* chore: fix lint

* fix: remove invalid test
---
 .../google/auth/oauth2/JwtCredentials.java    |  4 +-
 .../auth/oauth2/JwtCredentialsTest.java       | 50 +++++++++++++------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/oauth2_http/java/com/google/auth/oauth2/JwtCredentials.java b/oauth2_http/java/com/google/auth/oauth2/JwtCredentials.java
index 70140d65e..55fa3e5fb 100644
--- a/oauth2_http/java/com/google/auth/oauth2/JwtCredentials.java
+++ b/oauth2_http/java/com/google/auth/oauth2/JwtCredentials.java
@@ -86,7 +86,7 @@ public class JwtCredentials extends Credentials implements JwtProvider {
 
   private JwtCredentials(Builder builder) {
     this.privateKey = Preconditions.checkNotNull(builder.getPrivateKey());
-    this.privateKeyId = Preconditions.checkNotNull(builder.getPrivateKeyId());
+    this.privateKeyId = builder.getPrivateKeyId();
     this.jwtClaims = Preconditions.checkNotNull(builder.getJwtClaims());
     Preconditions.checkState(jwtClaims.isComplete(), JWT_INCOMPLETE_ERROR_MESSAGE);
     this.lifeSpanSeconds = Preconditions.checkNotNull(builder.getLifeSpanSeconds());
@@ -220,7 +220,7 @@ public PrivateKey getPrivateKey() {
     }
 
     public Builder setPrivateKeyId(String privateKeyId) {
-      this.privateKeyId = Preconditions.checkNotNull(privateKeyId);
+      this.privateKeyId = privateKeyId;
       return this;
     }
 
diff --git a/oauth2_http/javatests/com/google/auth/oauth2/JwtCredentialsTest.java b/oauth2_http/javatests/com/google/auth/oauth2/JwtCredentialsTest.java
index 44bad4e61..b7f328951 100644
--- a/oauth2_http/javatests/com/google/auth/oauth2/JwtCredentialsTest.java
+++ b/oauth2_http/javatests/com/google/auth/oauth2/JwtCredentialsTest.java
@@ -114,22 +114,6 @@ public void builder_requiresPrivateKey() {
     }
   }
 
-  @Test
-  public void builder_requiresPrivateKeyId() {
-    try {
-      JwtClaims claims =
-          JwtClaims.newBuilder()
-              .setAudience("some-audience")
-              .setIssuer("some-issuer")
-              .setSubject("some-subject")
-              .build();
-      JwtCredentials.newBuilder().setJwtClaims(claims).setPrivateKey(getPrivateKey()).build();
-      fail("Should throw exception");
-    } catch (NullPointerException ex) {
-      // expected
-    }
-  }
-
   @Test
   public void builder_requiresClaims() {
     try {
@@ -248,6 +232,40 @@ public void getRequestMetadata_withAdditionalClaims_hasJwtAccess() throws IOExce
         Collections.singletonMap("foo", "bar"));
   }
 
+  @Test
+  public void privateKeyIdNull() throws IOException {
+    JwtClaims claims =
+        JwtClaims.newBuilder()
+            .setAudience("some-audience")
+            .setIssuer("some-issuer")
+            .setSubject("some-subject")
+            .build();
+    JwtCredentials credentials =
+        JwtCredentials.newBuilder()
+            .setJwtClaims(claims)
+            .setPrivateKey(getPrivateKey())
+            .setPrivateKeyId(null)
+            .build();
+
+    Map<String, List<String>> metadata = credentials.getRequestMetadata();
+    verifyJwtAccess(metadata, "some-audience", "some-issuer", "some-subject", null);
+  }
+
+  @Test
+  public void privateKeyIdNotSpecified() throws IOException {
+    JwtClaims claims =
+        JwtClaims.newBuilder()
+            .setAudience("some-audience")
+            .setIssuer("some-issuer")
+            .setSubject("some-subject")
+            .build();
+    JwtCredentials credentials =
+        JwtCredentials.newBuilder().setJwtClaims(claims).setPrivateKey(getPrivateKey()).build();
+
+    Map<String, List<String>> metadata = credentials.getRequestMetadata();
+    verifyJwtAccess(metadata, "some-audience", "some-issuer", "some-subject", null);
+  }
+
   private void verifyJwtAccess(
       Map<String, List<String>> metadata,
       String expectedAudience,