Adding JWT payload compression

build.gradle

@@ -153,6 +153,8 @@ dependencies {
     implementation group: 'software.amazon.msk', name: 'aws-msk-iam-auth', version: '2.0.0'
 
     implementation group: 'io.projectreactor', name: 'reactor-core', version: '3.5.11'
+
+    implementation 'io.jsonwebtoken:jjwt-impl:0.12.3'
 }
 
 /**********************************************************************************************************************\

src/main/java/org/akhq/controllers/AbstractController.java

@@ -11,14 +11,12 @@
 import org.akhq.configs.security.Group;
 import org.akhq.configs.security.SecurityProperties;
 import org.akhq.security.annotation.AKHQSecured;
+import org.akhq.security.rule.AKHQSecurityRule;
 
 import java.lang.reflect.Method;
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -52,8 +50,7 @@ protected List<Group> getUserGroups() {
             return List.of();
         }
 
-        List<Group> groupBindings = ((Map<String, List<?>>) authentication.get().getAttributes().get("groups"))
-            .values()
+        List<Group> groupBindings = AKHQSecurityRule.decompressGroups(authentication.get()).values()
             .stream()
             .flatMap(Collection::stream)
             .map(gb -> new ObjectMapper().convertValue(gb, Group.class))
@@ -155,7 +152,7 @@ protected void checkIfClusterAndResourceAllowed(String cluster, String resource)
         try {
             AKHQSecured annotation = getCallingAKHQSecuredAnnotation();
 
-            isAllowed = ((Map<String, List<?>>)authentication.get().getAttributes().get("groups")).values().stream()
+            isAllowed = AKHQSecurityRule.decompressGroups(authentication.get()).values().stream()
                 .flatMap(Collection::stream)
                 .map(gb -> new ObjectMapper().convertValue(gb, Group.class))
                 // Get only group with role matching the method annotation resource and action

src/main/java/org/akhq/security/claim/CustomClaimsGenerator.java

@@ -0,0 +1,45 @@
+package org.akhq.security.claim;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.nimbusds.jwt.JWTClaimsSet;
+import io.jsonwebtoken.impl.compression.GzipCompressionAlgorithm;
+import io.micronaut.context.annotation.Replaces;
+import io.micronaut.core.annotation.Nullable;
+import io.micronaut.runtime.ApplicationConfiguration;
+import io.micronaut.security.authentication.Authentication;
+import io.micronaut.security.token.claims.ClaimsAudienceProvider;
+import io.micronaut.security.token.claims.JtiGenerator;
+import io.micronaut.security.token.config.TokenConfiguration;
+import io.micronaut.security.token.jwt.generator.claims.JWTClaimsSetGenerator;
+import jakarta.inject.Singleton;
+
+import java.util.Base64;
+
+@Singleton
+@Replaces(JWTClaimsSetGenerator.class)
+public class CustomClaimsGenerator extends JWTClaimsSetGenerator {
+    private final GzipCompressionAlgorithm gzip = new GzipCompressionAlgorithm();
+    private final ObjectMapper mapper = new ObjectMapper();
+
+    /**
+     * @param tokenConfiguration       Token Configuration
+     * @param jwtIdGenerator           Generator which creates unique JWT ID
+     * @param claimsAudienceProvider   Provider which identifies the recipients that the JWT is intended for.
+     * @param applicationConfiguration The application configuration
+     */
+    public CustomClaimsGenerator(TokenConfiguration tokenConfiguration, @Nullable JtiGenerator jwtIdGenerator, @Nullable ClaimsAudienceProvider claimsAudienceProvider, @Nullable ApplicationConfiguration applicationConfiguration) {
+        super(tokenConfiguration, jwtIdGenerator, claimsAudienceProvider, applicationConfiguration);
+    }
+
+    protected void populateWithAuthentication(JWTClaimsSet.Builder builder, Authentication authentication) {
+        super.populateWithAuthentication(builder, authentication);
+
+        try {
+            String plainGroups = mapper.writeValueAsString(builder.getClaims().get("groups"));
+            builder.claim("groups", Base64.getEncoder().encodeToString(gzip.compress(plainGroups.getBytes())));
+        } catch (JsonProcessingException e) {
+            throw new RuntimeException("Unable to build the JWT token, groups format is incorrect");
+        }
+    }
+}

src/main/java/org/akhq/security/rule/AKHQSecurityRule.java

@@ -1,6 +1,7 @@
 package org.akhq.security.rule;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import io.jsonwebtoken.impl.compression.GzipCompressionAlgorithm;
 import io.micronaut.http.HttpAttributes;
 import io.micronaut.http.HttpRequest;
 import io.micronaut.security.authentication.Authentication;
@@ -20,16 +21,16 @@
 import org.akhq.security.annotation.AKHQSecured;
 import org.reactivestreams.Publisher;
 
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
+import java.util.*;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 @Slf4j
 @Singleton
 public class AKHQSecurityRule extends AbstractSecurityRule<HttpRequest<?>> {
+    private static final ObjectMapper mapper = new ObjectMapper();
+    private static final GzipCompressionAlgorithm gzip = new GzipCompressionAlgorithm();
+
     /**
      * @param rolesFinder Roles Parser
      */
@@ -69,7 +70,7 @@ public Publisher<SecurityRuleResult> check(HttpRequest<?> request, Authenticatio
             return Flowable.just(SecurityRuleResult.REJECTED);
         }
 
-        List<Group> userGroups = ((Map<String, List<?>>)authentication.getAttributes().get("groups")).values().stream()
+        List<Group> userGroups = decompressGroups(authentication).values().stream()
             .flatMap(Collection::stream)
             // Type mismatch during serialization from LinkedTreeMap to Group if we use List<Group>
             // Need to serialize Object to Group manually in the stream
@@ -98,6 +99,17 @@ public Publisher<SecurityRuleResult> check(HttpRequest<?> request, Authenticatio
             return Flowable.just(SecurityRuleResult.REJECTED);
     }
 
+    public static Map<String, List<?>> decompressGroups(Authentication authentication) {
+        try {
+            String base64CompressedGroups = ((String) authentication.getAttributes().get("groups"));
+            String compressedGroups = new String(gzip.decompress(Base64.getDecoder().decode(base64CompressedGroups)));
+            return mapper.readValue(compressedGroups, Map.class);
+        } catch (Exception e) {
+            log.trace("JWT payload is not compressed, returning groups directly");
+            return (Map<String, List<?>>) authentication.getAttributes().get("groups");
+        }
+    }
+
     public static final Integer ORDER = SecuredAnnotationRule.ORDER - 100;
 
     public int getOrder() {

src/test/java/org/akhq/security/claim/RestApiClaimProviderTest.java

@@ -1,8 +1,10 @@
 package org.akhq.security.claim;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.nimbusds.jwt.JWT;
 import com.nimbusds.jwt.JWTParser;
+import io.jsonwebtoken.impl.compression.GzipCompressionAlgorithm;
 import io.micronaut.context.annotation.Replaces;
 import io.micronaut.context.annotation.Requires;
 import io.micronaut.core.util.StringUtils;
@@ -28,6 +30,7 @@
 import org.reactivestreams.Publisher;
 
 import java.text.ParseException;
+import java.util.Base64;
 import java.util.List;
 import java.util.Map;
 import jakarta.annotation.security.PermitAll;
@@ -45,7 +48,7 @@ public class RestApiClaimProviderTest {
     protected RxHttpClient client;
 
     @Test
-    void loginExternalClaim() throws ParseException {
+    void loginExternalClaim() throws ParseException, JsonProcessingException {
         HttpResponse<?> resultResponse = client.toBlocking().exchange(
             HttpRequest.POST("/login", new UsernamePasswordCredentials("admin", "pass"))
         );
@@ -57,7 +60,11 @@ void loginExternalClaim() throws ParseException {
 
         assertTrue(token.getJWTClaimsSet().getClaims().containsKey("groups"));
 
-        Map<String, List<Group>> groups = (Map<String, List<Group>>) token.getJWTClaimsSet().getClaim("groups");
+        var gzip = new GzipCompressionAlgorithm();
+        String compressedGroups = new String(gzip.decompress(
+            Base64.getDecoder().decode((String) token.getJWTClaimsSet().getClaim("groups"))));
+
+        Map<String, List<Group>> groups = new ObjectMapper().readValue(compressedGroups, Map.class);
         assertThat(groups.keySet(), hasSize(1));
         assertNotNull(groups.get("limited"));