usage jwt as ratelimit key

membrane · Aug 23, 2023 · 4776d11 · 4776d11
1 parent 5e19e1c
commit 4776d11
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 1 deletion.
diff --git a/core/src/main/java/com/predic8/membrane/core/lang/spel/ExchangeEvaluationContext.java b/core/src/main/java/com/predic8/membrane/core/lang/spel/ExchangeEvaluationContext.java
@@ -15,12 +15,15 @@
 package com.predic8.membrane.core.lang.spel;
 
 import com.fasterxml.jackson.databind.*;
+import com.predic8.membrane.annot.model.doc.Doc;
 import com.predic8.membrane.core.exchange.*;
 import com.predic8.membrane.core.http.*;
+import org.jose4j.jwt.JwtClaims;
 import org.springframework.expression.spel.support.*;
 
 import java.io.*;
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class ExchangeEvaluationContext {
 
@@ -36,7 +39,15 @@ public class ExchangeEvaluationContext {
     public ExchangeEvaluationContext(Exchange exchange, Message message) {
         this.exchange = exchange;
         this.message = message;
-        properties = exchange.getProperties();
+        properties = exchange.getProperties().entrySet().stream()
+                .map(entry -> {
+                    if (entry.getKey().equals("jwt") && entry.getValue() instanceof JwtClaims) {
+                        return new AbstractMap.SimpleEntry<String, Object>("jwt", ((JwtClaims) entry.getValue()).getClaimsMap());
+                    }
+                    return entry;
+                })
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
+        ;
         this.headers = new HeaderMap(message.getHeader());
 
         Request request = exchange.getRequest();

diff --git a/...c/test/java/com/predic8/membrane/core/interceptor/ratelimit/RateLimitInterceptorTest.java b/...c/test/java/com/predic8/membrane/core/interceptor/ratelimit/RateLimitInterceptorTest.java
@@ -21,13 +21,15 @@
 import com.predic8.membrane.core.interceptor.*;
 import com.predic8.membrane.core.util.*;
 import org.jetbrains.annotations.*;
+import org.jose4j.jwt.JwtClaims;
 import org.junit.jupiter.api.*;
 import org.junit.jupiter.params.*;
 import org.junit.jupiter.params.provider.*;
 
 import java.net.*;
 import java.util.*;
 import java.util.concurrent.atomic.*;
+import java.util.stream.IntStream;
 
 import static com.predic8.membrane.core.http.Header.*;
 import static com.predic8.membrane.core.interceptor.Outcome.*;
@@ -120,6 +122,34 @@ void testHandleRequestRateLimit1Second() throws Exception {
 		assertEquals(RETURN, rli.handleRequest(exc));
 
 	}
+
+	@Test
+	void rateLimitByJWT() throws Exception {
+		var interceptor = new RateLimitInterceptor(ofSeconds(10), 100);
+		interceptor.setKeyExpression("properties[jwt][sub]");
+		interceptor.init();
+
+		var exc = new Request.Builder().buildExchange();
+
+		// done by JwtAuthInterceptor
+		var claims = new JwtClaims();
+		claims.setSubject("fooman");
+		exc.getProperties().put("jwt", claims);
+
+		IntStream.range(0, interceptor.getRequestLimit())
+				.parallel()
+				.forEach(i -> {
+					try {
+						var outcome = interceptor.handleRequest(exc);
+						assertEquals(CONTINUE, outcome);
+					} catch (Exception e) {
+						throw new RuntimeException(e);
+					}
+				});
+
+		var outcome = interceptor.handleRequest(exc);
+		assertEquals(RETURN, outcome);
+	}
 
 	@Test
 	void handleRequestRateLimit1SecondConcurrency() throws Exception