Fixing JwtCreator builder when setting headers as a map

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;
         }
 

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 shouldOverwriteExistingIfHeadersMapContainsTheSameKey() 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()

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 + "\":";
     }