Allow removing claims when bulding JWT (#710)

smallrye · Jul 5, 2023 · aaa5f8c · aaa5f8c
1 parent 4a49a2d
commit aaa5f8c
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 0 deletions.
diff --git a/implementation/jwt-build/src/main/java/io/smallrye/jwt/build/JwtClaimsBuilder.java b/implementation/jwt-build/src/main/java/io/smallrye/jwt/build/JwtClaimsBuilder.java
@@ -204,6 +204,14 @@ default JwtClaimsBuilder claim(Claims name, Object value) {
      */
     JwtClaimsBuilder claim(String name, Object value);
 
+    /**
+     * Remove a claim.
+     *
+     * @param name the claim name
+     * @return JwtClaimsBuilder
+     */
+    JwtClaimsBuilder remove(String name);
+
     /**
      * Set JsonWebSignature headers and sign the claims by moving to {@link JwtSignatureBuilder}
      *

diff --git a/implementation/jwt-build/src/main/java/io/smallrye/jwt/build/impl/JwtClaimsBuilderImpl.java b/implementation/jwt-build/src/main/java/io/smallrye/jwt/build/impl/JwtClaimsBuilderImpl.java
@@ -378,4 +378,10 @@ public Object verify(String name, Object value) {
         }
     }
 
+    @Override
+    public JwtClaimsBuilder remove(String name) {
+        claims.unsetClaim(name);
+        return this;
+    }
+
 }
diff --git a/implementation/jwt-build/src/test/java/io/smallrye/jwt/build/JwtSignTest.java b/implementation/jwt-build/src/test/java/io/smallrye/jwt/build/JwtSignTest.java
@@ -19,6 +19,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
@@ -39,6 +40,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.UUID;
 
 import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
@@ -110,6 +112,30 @@ void enhanceAndResignToken() throws Exception {
         assertEquals("https://localhost:8081", claims.getAudience().get(0));
     }
 
+    @Test
+    void enhanceAndResignTokenWithCustomClaimRemoved() throws Exception {
+        JwtClaims tokenClaims = signAndVerifyClaims();
+        assertEquals("custom-value", tokenClaims.getClaimValue("customClaim"));
+        JsonWebToken token = new TestJsonWebToken(tokenClaims);
+
+        String jwt = Jwt.claims(token).remove("customClaim")
+                // this just checks trying to remove non-existent claims does not cause some NPE
+                .remove(UUID.randomUUID().toString())
+                .claim("newClaim", "new-value").sign();
+
+        // verify
+        JsonWebSignature jws = getVerifiedJws(jwt);
+        JwtClaims claims = JwtClaims.parse(jws.getPayload());
+        assertEquals(6, claims.getClaimsMap().size());
+        checkDefaultClaimsAndHeaders(getJwsHeaders(jwt, 2), claims);
+        assertNull(claims.getClaimValue("customClaim"));
+
+        assertEquals("new-value", claims.getClaimValue("newClaim"));
+        assertEquals("https://default-issuer", claims.getIssuer());
+        assertEquals(1, claims.getAudience().size());
+        assertEquals("https://localhost:8081", claims.getAudience().get(0));
+    }
+
     @Test
     void enhanceAndResignTokenWithConfiguredIssuerAndAudUsed() throws Exception {
         JsonWebToken token = new TestJsonWebToken(signAndVerifyClaims());