Merge pull request #320 from maxbalan/fixing_jwt_builder_headers

Fixing JwtCreator builder when setting headers as a map
auth0 · Nov 6, 2019 · 2a0c022 · 2a0c022
2 parents 9d69682 + bb83717
commit 2a0c022
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 2 deletions.
diff --git a/lib/src/main/java/com/auth0/jwt/JWTCreator.java b/lib/src/main/java/com/auth0/jwt/JWTCreator.java
@@ -66,12 +66,25 @@ public static class Builder {
 
         /**
          * Add specific Claims to set as the Header.
+         * If provided map is null then nothing is changed
+         * If provided map contains a claim with null value then that claim will be removed from the header
          *
          * @param headerClaims the values to use as Claims in the token's Header.
          * @return this same Builder instance.
          */
         public Builder withHeader(Map<String, Object> headerClaims) {
-            this.headerClaims = new HashMap<>(headerClaims);
+            if (headerClaims == null) {
+                return this;
+            }
+
+            for (Map.Entry<String, Object> entry : headerClaims.entrySet()) {
+                if (entry.getValue() == null) {
+                    this.headerClaims.remove(entry.getKey());
+                } else {
+                    this.headerClaims.put(entry.getKey(), entry.getValue());
+                }
+            }
+
             return this;
         }
 

diff --git a/lib/src/test/java/com/auth0/jwt/JWTCreatorTest.java b/lib/src/test/java/com/auth0/jwt/JWTCreatorTest.java
@@ -1,6 +1,7 @@
 package com.auth0.jwt;
 
 import com.auth0.jwt.algorithms.Algorithm;
+import com.auth0.jwt.impl.PublicClaims;
 import com.auth0.jwt.interfaces.ECDSAKeyProvider;
 import com.auth0.jwt.interfaces.RSAKeyProvider;
 import org.apache.commons.codec.binary.Base64;
@@ -53,6 +54,65 @@ public void shouldAddHeaderClaim() throws Exception {
         assertThat(headerJson, JsonMatcher.hasEntry("asd", 123));
     }
 
+    @Test
+    public void shouldReturnBuilderIfNullMapIsProvided() throws Exception {
+        String signed = JWTCreator.init()
+                                  .withHeader(null)
+                                  .sign(Algorithm.HMAC256("secret"));
+
+        assertThat(signed, is(notNullValue()));
+    }
+
+    @Test
+    public void shouldOverwriteExistingHeaderIfHeaderMapContainsTheSameKey() throws Exception {
+        Map<String, Object> header = new HashMap<String, Object>();
+        header.put(PublicClaims.KEY_ID, "xyz");
+
+        String signed = JWTCreator.init()
+                                  .withKeyId("abc")
+                                  .withHeader(header)
+                                  .sign(Algorithm.HMAC256("secret"));
+
+        assertThat(signed, is(notNullValue()));
+        String[] parts = signed.split("\\.");
+        String headerJson = new String(Base64.decodeBase64(parts[0]), StandardCharsets.UTF_8);
+        assertThat(headerJson, JsonMatcher.hasEntry(PublicClaims.KEY_ID, "xyz"));
+    }
+
+    @Test
+    public void shouldOverwriteExistingHeadersWhenSettingSameHeaderKey() throws Exception {
+        Map<String, Object> header = new HashMap<String, Object>();
+        header.put(PublicClaims.KEY_ID, "xyz");
+
+        String signed = JWTCreator.init()
+                                  .withHeader(header)
+                                  .withKeyId("abc")
+                                  .sign(Algorithm.HMAC256("secret"));
+
+        assertThat(signed, is(notNullValue()));
+        String[] parts = signed.split("\\.");
+        String headerJson = new String(Base64.decodeBase64(parts[0]), StandardCharsets.UTF_8);
+        assertThat(headerJson, JsonMatcher.hasEntry(PublicClaims.KEY_ID, "abc"));
+    }
+
+    @Test
+    public void shouldRemoveHeaderIfTheValueIsNull() throws Exception {
+        Map<String, Object> header = new HashMap<String, Object>();
+        header.put(PublicClaims.KEY_ID, null);
+        header.put("test2", "isSet");
+
+        String signed = JWTCreator.init()
+                                  .withKeyId("test")
+                                  .withHeader(header)
+                                  .sign(Algorithm.HMAC256("secret"));
+
+        assertThat(signed, is(notNullValue()));
+        String[] parts = signed.split("\\.");
+        String headerJson = new String(Base64.decodeBase64(parts[0]), StandardCharsets.UTF_8);
+        assertThat(headerJson, JsonMatcher.isNotPresent(PublicClaims.KEY_ID));
+        assertThat(headerJson, JsonMatcher.hasEntry("test2", "isSet"));
+    }
+
     @Test
     public void shouldAddKeyId() throws Exception {
         String signed = JWTCreator.init()
@@ -357,4 +417,4 @@ public void shouldAcceptCustomArrayClaimOfTypeLong() throws Exception {
         String[] parts = jwt.split("\\.");
         assertThat(parts[1], is("eyJuYW1lIjpbMSwyLDNdfQ"));
     }
-}
+}
diff --git a/lib/src/test/java/com/auth0/jwt/JsonMatcher.java b/lib/src/test/java/com/auth0/jwt/JsonMatcher.java
@@ -68,6 +68,10 @@ public static JsonMatcher hasEntry(String key, Matcher valueMatcher) {
         return new JsonMatcher(key, null, valueMatcher);
     }
 
+    public static JsonMatcher isNotPresent(String key) {
+        return new JsonMatcher(key, null, null);
+    }
+
     private String getStringKey(String key) {
         return "\"" + key + "\":";
     }