* sse: removed default rate limit group ServerSentEventConfig.SSE_CONNECT_GROUP ("sse:connect")

* sse: support @limitrate on sse connect or class

Signed-off-by: neo <1100909+neowu@users.noreply.github.com>

Author: neowu

CHANGELOG.md

@@ -1,5 +1,11 @@
 ## Change log
 
+### 9.1.8 (3/7/2025 - )
+
+* sse: removed default rate limit group ServerSentEventConfig.SSE_CONNECT_GROUP ("sse:connect")
+* sse: support @LimitRate on sse connect or class
+  > use @LimitRate and http().limitRate().add() to configure rate limit
+
 ### 9.1.7 (2/26/2025 - 3/6/2025)
 
 * sse: send ErrorResponse to client via "event: error" on exception

build.gradle.kts

@@ -7,7 +7,7 @@ apply(plugin = "project")
 
 subprojects {
     group = "core.framework"
-    version = "9.1.7"
+    version = "9.1.8"
 }
 
 val elasticVersion = "8.15.0"

core-ng/src/main/java/core/framework/internal/web/HTTPHandlerContext.java

@@ -13,6 +13,7 @@ public class HTTPHandlerContext {
     public final RequestBeanReader requestBeanReader = new RequestBeanReader();
     public final ResponseBeanWriter responseBeanWriter = new ResponseBeanWriter();
     public final RateControl rateControl = new RateControl();
+    public boolean limitRate;
     @Nullable
     public IPv4AccessControl accessControl;
 }

core-ng/src/main/java/core/framework/internal/web/HTTPIOHandler.java

@@ -70,7 +70,7 @@ public void handleRequest(HttpServerExchange exchange) throws Exception {
         HttpString method = exchange.getRequestMethod();
         HeaderMap headers = exchange.getRequestHeaders();
 
-        var requestHandler = new HTTPRequestHandler(exchange, handler, sseHandler);
+        var requestHandler = new Handler(exchange);
         boolean ws = webSocketHandler != null && webSocketHandler.check(method, headers);   // TODO: retire ws and simplify
         boolean active = !requestHandler.sse && !ws;
         boolean shutdown = shutdownHandler.handle(exchange, active);
@@ -123,4 +123,26 @@ boolean hasBody(long contentLength, HttpString method) {
         if (contentLength == 0) return false;  // if body is empty, skip reading
         return Methods.POST.equals(method) || Methods.PUT.equals(method) || Methods.PATCH.equals(method);
     }
+
+    public class Handler {
+        private final boolean sse;
+        private final HttpServerExchange exchange;
+
+        Handler(HttpServerExchange exchange) {
+            this.exchange = exchange;
+
+            HttpString method = exchange.getRequestMethod();
+            String path = exchange.getRequestPath();
+            HeaderMap headers = exchange.getRequestHeaders();
+            sse = sseHandler != null && sseHandler.check(method, path, headers);
+        }
+
+        public void handle() {
+            if (sse) {
+                sseHandler.handleRequest(exchange); // not dispatch, continue in io thread
+            } else {
+                exchange.dispatch(handler);
+            }
+        }
+    }
 }

core-ng/src/main/java/core/framework/internal/web/HTTPRequestHandler.java

@@ -1,6 +1,6 @@
 package core.framework.internal.web.request;
 
-import core.framework.internal.web.HTTPRequestHandler;
+import core.framework.internal.web.HTTPIOHandler;
 import core.framework.web.exception.BadRequestException;
 import io.undertow.connector.PooledByteBuffer;
 import io.undertow.server.HttpServerExchange;
@@ -20,13 +20,13 @@ public final class RequestBodyReader implements ChannelListener<StreamSourceChan
     static final AttachmentKey<RequestBody> REQUEST_BODY = AttachmentKey.create(RequestBody.class);
 
     private final HttpServerExchange exchange;
-    private final HTTPRequestHandler handler;
+    private final HTTPIOHandler.Handler handler;
     private final int contentLength;
     private boolean complete;
     private byte[] body;
     private int position = 0;
 
-    public RequestBodyReader(HttpServerExchange exchange, HTTPRequestHandler handler) {
+    public RequestBodyReader(HttpServerExchange exchange, HTTPIOHandler.Handler handler) {
         this.exchange = exchange;
         this.handler = handler;
         contentLength = (int) exchange.getRequestContentLength();

core-ng/src/main/java/core/framework/internal/web/request/RequestBodyReader.java

@@ -1,15 +1,33 @@
 package core.framework.internal.web.sse;
 
+import core.framework.web.Request;
+import core.framework.web.sse.Channel;
 import core.framework.web.sse.ChannelListener;
 
+import java.lang.annotation.Annotation;
+import java.lang.reflect.Method;
+
 class ChannelSupport<T> {
     final ChannelListener<T> listener;
     final ServerSentEventContextImpl<T> context;
     final ServerSentEventBuilder<T> builder;
+    final Method targetMethod;
 
     ChannelSupport(ChannelListener<T> listener, Class<T> eventClass, ServerSentEventContextImpl<T> context) {
         this.listener = listener;
         this.context = context;
         builder = new ServerSentEventBuilder<>(eventClass);
+        try {
+            this.targetMethod = listener.getClass().getMethod("onConnect", Request.class, Channel.class, String.class);
+        } catch (NoSuchMethodException e) {
+            throw new Error("failed to get listener.onConnect method", e);
+        }
+    }
+
+    <V extends Annotation> V annotation(Class<V> annotationClass) {
+        V annotation = targetMethod.getDeclaredAnnotation(annotationClass);
+        if (annotation == null)
+            annotation = listener.getClass().getDeclaredAnnotation(annotationClass);
+        return annotation;
     }
 }

core-ng/src/main/java/core/framework/internal/web/sse/ChannelSupport.java

@@ -5,12 +5,13 @@
 import core.framework.internal.log.ActionLog;
 import core.framework.internal.log.LogManager;
 import core.framework.internal.web.HTTPHandlerContext;
+import core.framework.internal.web.http.RateControl;
 import core.framework.internal.web.request.RequestImpl;
 import core.framework.internal.web.service.ErrorResponse;
 import core.framework.internal.web.session.ReadOnlySession;
 import core.framework.internal.web.session.SessionManager;
-import core.framework.module.ServerSentEventConfig;
 import core.framework.util.Strings;
+import core.framework.web.rate.LimitRate;
 import core.framework.web.sse.ChannelListener;
 import io.undertow.server.HttpHandler;
 import io.undertow.server.HttpServerExchange;
@@ -95,7 +96,10 @@ void handle(HttpServerExchange exchange, StreamSinkChannel sink) {
             @SuppressWarnings("unchecked")
             ChannelSupport<Object> support = (ChannelSupport<Object>) supports.get(key(request.method().name(), path));   // ServerSentEventHandler.check() ensures path exists
             actionLog.action("sse:" + path + ":connect");
-            handlerContext.rateControl.validateRate(ServerSentEventConfig.SSE_CONNECT_GROUP, request.clientIP());
+
+            if (handlerContext.limitRate) {
+                limitRate(handlerContext.rateControl, support, request.clientIP());
+            }
 
             channel = new ChannelImpl<>(exchange, sink, support.context, support.builder, actionLog.id);
             actionLog.context("channel", channel.id);
@@ -132,6 +136,14 @@ void handle(HttpServerExchange exchange, StreamSinkChannel sink) {
         }
     }
 
+    void limitRate(RateControl rateControl, ChannelSupport<Object> support, String clientIP) {
+        LimitRate limitRate = support.annotation(LimitRate.class);
+        if (limitRate != null) {
+            String group = limitRate.value();
+            rateControl.validateRate(group, clientIP);
+        }
+    }
+
     byte[] errorResponse(byte[] errorResponse) {
         ByteBuffer buffer = ByteBuffer.wrap(new byte[errorResponse.length + 38]);
         buffer.put(Strings.bytes("retry: 86400000\n\nevent: error\ndata: "));   // tell browser retry in 24 hours

core-ng/src/main/java/core/framework/internal/web/sse/ServerSentEventHandler.java

@@ -20,6 +20,7 @@ protected void initialize(ModuleContext context, String name) {
         // save at max 5K group/ip combination per pod, about 800K memory, to adapt with more ips/cc attack, better defense with cloud infra based solution together
         maxEntries(5000);
         context.httpServerConfig.interceptors.add(new LimitRateInterceptor(rateControl));
+        context.httpServer.handlerContext.limitRate = true;
     }
 
     @Override

core-ng/src/main/java/core/framework/module/LimitRateConfig.java

@@ -18,9 +18,6 @@
 import java.time.Duration;
 
 public class ServerSentEventConfig extends Config {
-    // use http().limitRate().add(ServerSentEventConfig.SSE_CONNECT_GROUP, ...) to configure rate limiting for sse connections
-    public static final String SSE_CONNECT_GROUP = "sse:connect";
-
     private final Logger logger = LoggerFactory.getLogger(ServerSentEventConfig.class);
 
     ModuleContext context;
@@ -31,13 +28,6 @@ protected void initialize(ModuleContext context, @Nullable String name) {
         this.context = context;
     }
 
-    @Override
-    protected void validate() {
-        if (!context.httpServer.handlerContext.rateControl.hasGroup(SSE_CONNECT_GROUP)) {
-            context.httpServer.handlerContext.rateControl.config(SSE_CONNECT_GROUP, 10, 10, Duration.ofSeconds(30));
-        }
-    }
-
     public <T> void listen(HTTPMethod method, String path, Class<T> eventClass, ChannelListener<T> listener) {
         if (HTTPIOHandler.HEALTH_CHECK_PATH.equals(path)) throw new Error("/health-check is reserved path");
         if (path.contains("/:")) throw new Error("listener path must be static, path=" + path);