#190 Added request limiting

mangooio-core/src/main/java/io/mangoo/core/Bootstrap.java

@@ -159,14 +159,15 @@ public void parseRoutes() {
 
             if (yamlRouter != null) {
                 for (final YamlRoute yamlRoute : yamlRouter.getRoutes()) {
-                    final Route route = new Route(BootstrapUtils.getRouteType(yamlRoute.getMethod()));
-                    route.toUrl(yamlRoute.getUrl().trim());
-                    route.withRequest(HttpString.tryFromString(yamlRoute.getMethod()));
-                    route.withUsername(yamlRoute.getUsername());
-                    route.withPassword(yamlRoute.getPassword());
-                    route.withAuthentication(yamlRoute.isAuthentication());
-                    route.withTimer(yamlRoute.isTimer());
-                    route.allowBlocking(yamlRoute.isBlocking());
+                    final Route route = new Route(BootstrapUtils.getRouteType(yamlRoute.getMethod()))
+                            .toUrl(yamlRoute.getUrl().trim())
+                            .withRequest(HttpString.tryFromString(yamlRoute.getMethod()))
+                            .withUsername(yamlRoute.getUsername())
+                            .withPassword(yamlRoute.getPassword())
+                            .withAuthentication(yamlRoute.isAuthentication())
+                            .withTimer(yamlRoute.isTimer())
+                            .withLimit(yamlRoute.getLimit())
+                            .allowBlocking(yamlRoute.isBlocking());
 
                     String mapping = yamlRoute.getMapping();   
                     try {
@@ -250,7 +251,8 @@ private RoutingHandler getRoutingHandler() {
                         route.isInternalTemplateEngine(),
                         route.isTimerEnabled(),
                         route.getUsername(),
-                        route.getPassword());
+                        route.getPassword(),
+                        route.getLimit());
 
                 routingHandler.add(route.getRequestMethod(),route.getUrl(), dispatcherHandler);
             } else if (RouteType.RESOURCE_FILE.equals(route.getRouteType())) {
@@ -399,6 +401,7 @@ public static class YamlRoute {
         private String mapping;
         private String username;
         private String password;
+        private int limit;
         private boolean blocking;
         private boolean authentication;
         private boolean timer;
@@ -427,6 +430,10 @@ public void setUsername(String username) {
             this.username = username;
         } 
 
+        public void setLimit(int limit) {
+            this.limit = limit;
+        } 
+
         public String getPassword() {
             return password;
         }
@@ -439,6 +446,10 @@ public String getMapping() {
             return mapping;
         }
 
+        public int getLimit() {
+            return limit;
+        }
+
         public void setMapping(String mapping) {
             this.mapping = mapping;
         }

mangooio-core/src/main/java/io/mangoo/routing/Attachment.java

@@ -6,6 +6,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import org.apache.commons.lang3.StringUtils;
+
 import io.mangoo.crypto.Crypto;
 import io.mangoo.i18n.Messages;
 import io.mangoo.routing.bindings.Authentication;
@@ -23,12 +25,15 @@
 public class Attachment {
     private final long start = System.currentTimeMillis();
     private int methodParametersCount;
+    private int limit;
     private Class<?> controllerClass;
     private Object controllerInstance;
     private Map<String, Class<?>> methodParameters;
     private String controllerClassName;
     private String controllerMethodName;
     private String body;
+    private String username;
+    private String password;
     private Method method;
     private Authentication authentication;
     private Session session;
@@ -124,6 +129,21 @@ public Attachment withTimer(boolean timer) {
         return this;
     }
 
+    public Attachment withLimit(int limit) {
+        this.limit = limit;
+        return this;
+    }
+
+    public Attachment withUsername(String username) {
+        this.username = username;
+        return this;
+    }
+
+    public Attachment withPassword(String password) {
+        this.password = password;
+        return this;
+    }
+
     public Messages getMessages() {
         return this.messages;
     }
@@ -231,6 +251,26 @@ public Response getResponse() {
     public boolean isTimer() {
         return this.timer;
     }
+
+    public String getUsername() {
+        return this.username;
+    }
+
+    public String getPassword() {
+        return this.password;
+    }
+
+    public int getLimit() {
+        return this.limit;
+    }
+
+    public boolean hasAuthentication() {
+        return StringUtils.isNotBlank(this.username) && StringUtils.isNotBlank(this.password);
+    }
+
+    public boolean hasLimit() {
+        return this.limit > 0;
+    }
 
     public long getResponseTime() {
         return System.currentTimeMillis() - this.start;

mangooio-core/src/main/java/io/mangoo/routing/Route.java

@@ -11,13 +11,14 @@
  *
  */
 public class Route {
+    private final RouteType routeType;
     private Class<?> controllerClass;
     private String controllerMethod;
     private HttpString requestMethod;
     private String url;
     private String username;
     private String password;
-    private final RouteType routeType;
+    private int limit;
     private boolean authentication;
     private boolean blocking;
     private boolean timer;
@@ -92,10 +93,19 @@ public Route useInternalTemplateEngine() {
         return this;
     }
 
+    public Route withLimit(int limit) {
+        this.limit = limit;
+        return this;
+    }
+
     public String getUrl() {
         return this.url;
     }
 
+    public int getLimit() {
+        return this.limit;
+    }
+
     public String getUsername() {
         return this.username;
     }    

mangooio-core/src/main/java/io/mangoo/routing/handlers/DispatcherHandler.java

@@ -4,13 +4,11 @@
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
-import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -20,18 +18,10 @@
 import io.mangoo.crypto.Crypto;
 import io.mangoo.i18n.Messages;
 import io.mangoo.interfaces.MangooRequestFilter;
-import io.mangoo.models.Identity;
 import io.mangoo.routing.Attachment;
 import io.mangoo.routing.listeners.MetricsListener;
 import io.mangoo.templating.TemplateEngine;
 import io.mangoo.utils.RequestUtils;
-import io.undertow.security.api.AuthenticationMechanism;
-import io.undertow.security.api.AuthenticationMode;
-import io.undertow.security.handlers.AuthenticationCallHandler;
-import io.undertow.security.handlers.AuthenticationConstraintHandler;
-import io.undertow.security.handlers.AuthenticationMechanismsHandler;
-import io.undertow.security.handlers.SecurityInitialHandler;
-import io.undertow.security.impl.BasicAuthenticationMechanism;
 import io.undertow.server.HttpHandler;
 import io.undertow.server.HttpServerExchange;
 
@@ -65,12 +55,13 @@ public class DispatcherHandler implements HttpHandler {
     private final String controllerMethodName;
     private final String username;
     private final String password;    
+    private final int limit;
     private final int methodParametersCount;
     private final boolean async;
     private final boolean timer;
     private final boolean hasRequestFilter;
 
-    public DispatcherHandler(Class<?> controllerClass, String controllerMethod, boolean async, boolean internalTemplateEngine, boolean timer, String username, String password) {
+    public DispatcherHandler(Class<?> controllerClass, String controllerMethod, boolean async, boolean internalTemplateEngine, boolean timer, String username, String password, int limit) {
         Objects.requireNonNull(controllerClass, "controllerClass can not be null");
         Objects.requireNonNull(controllerMethod, "controllerMethod can not be null");
 
@@ -86,6 +77,7 @@ public DispatcherHandler(Class<?> controllerClass, String controllerMethod, bool
         this.methodParametersCount = this.methodParameters.size();
         this.async = async;
         this.timer = timer;
+        this.limit = limit;
         this.hasRequestFilter = Application.getInjector().getAllBindings().containsKey(com.google.inject.Key.get(MangooRequestFilter.class));
 
         try {
@@ -134,6 +126,9 @@ public void handleRequest(HttpServerExchange exchange) throws Exception {
             .withRequestParameter(RequestUtils.getRequestParameters(exchange))
             .withMessages(this.messages)
             .withTimer(this.timer)
+            .withLimit(this.limit)
+            .withUsername(this.username)
+            .withPassword(this.password)
             .withTemplateEngine(this.templateEngine)
             .withCrypto(this.crypto);
 
@@ -166,38 +161,6 @@ private Map<String, Class<?>> getMethodParameters() {
      */
     @SuppressWarnings("all")
     private void nextHandler(HttpServerExchange exchange) throws Exception {
-        if (requestHasAuthentication()) {
-            HttpHandler httpHandler = addSecurity(Application.getInstance(AuthenticationHandler.class));
-            httpHandler.handleRequest(exchange);
-        } else {
-            Application.getInstance(LocaleHandler.class).handleRequest(exchange);
-        }
-    }
-
-    /**
-     * Checks if the request requires basic authentication
-     * 
-     * @return True if username and password are not blank, false otherwise
-     */
-    private boolean requestHasAuthentication() {
-        return StringUtils.isNotBlank(this.username) && StringUtils.isNotBlank(this.password);
-    }
-
-    /**
-     * Adds a Wrapper to the handler when the request requires authentication
-     * 
-     * @param wrap The Handler to wrap
-     * @return A wrapped handler
-     */
-    private HttpHandler addSecurity(final HttpHandler wrap) {
-        HttpHandler handler = wrap;
-
-        final List<AuthenticationMechanism> mechanisms = Collections.<AuthenticationMechanism>singletonList(new BasicAuthenticationMechanism("Authentication required"));
-        handler = new AuthenticationCallHandler(handler);
-        handler = new AuthenticationConstraintHandler(handler);
-        handler = new AuthenticationMechanismsHandler(handler, mechanisms);
-        handler = new SecurityInitialHandler(AuthenticationMode.PRO_ACTIVE, new Identity(this.username, this.password), handler);
-
-        return handler;
+        Application.getInstance(LimitHandler.class).handleRequest(exchange);
     }
 }

mangooio-core/src/main/java/io/mangoo/routing/handlers/LimitHandler.java

@@ -0,0 +1,84 @@
+package io.mangoo.routing.handlers;
+
+import java.net.InetSocketAddress;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import com.google.inject.Inject;
+import com.tc.text.StringUtils;
+
+import io.mangoo.cache.Cache;
+import io.mangoo.core.Application;
+import io.mangoo.enums.CacheName;
+import io.mangoo.providers.CacheProvider;
+import io.mangoo.routing.Attachment;
+import io.mangoo.utils.RequestUtils;
+import io.undertow.server.HttpHandler;
+import io.undertow.server.HttpServerExchange;
+import io.undertow.util.StatusCodes;
+
+/**
+ * 
+ * @author svenkubiak
+ *
+ */
+public class LimitHandler implements HttpHandler {
+    private Attachment requestAttachment;
+    private Cache cache;
+
+    @Inject
+    public LimitHandler(CacheProvider cacheProvider) {
+        Objects.requireNonNull(cacheProvider, "cacheProvider can not be null");
+        this.cache = cacheProvider.getCache(CacheName.REQUEST);
+    }
+
+    @Override
+    public void handleRequest(HttpServerExchange exchange) throws Exception {
+        this.requestAttachment = exchange.getAttachment(RequestUtils.ATTACHMENT_KEY);
+        if (this.requestAttachment.hasLimit()) {
+            String url = exchange.getRequestURL();
+            InetSocketAddress inetSocketAddress = exchange.getSourceAddress();
+            if (StringUtils.isNotBlank(url) && inetSocketAddress != null) {
+                String hostString = inetSocketAddress.getHostString();
+                String key = hostString.trim().toLowerCase() + url.trim().toLowerCase();
+
+                AtomicInteger counter = this.cache.get(key) ;
+                if (counter == null) {
+                    counter = new AtomicInteger();
+                }
+
+                if (this.requestAttachment.getLimit() >= counter.get()) {
+                    counter.incrementAndGet();
+                    this.cache.put(key, counter);   
+                } else {
+                    endRequest(exchange); 
+                }
+            } else {
+                endRequest(exchange);
+            }
+        } else {
+            nextHandler(exchange);
+        }
+    }
+
+    private void endRequest(HttpServerExchange exchange) {
+        exchange.setStatusCode(StatusCodes.TOO_MANY_REQUESTS);
+        exchange.endExchange(); 
+    }
+
+    /**
+     * Handles the next request in the handler chain
+     *
+     * @param exchange The HttpServerExchange
+     * @throws Exception Thrown when an exception occurs
+     */
+    @SuppressWarnings("all")
+    protected void nextHandler(HttpServerExchange exchange) throws Exception {
+        if (this.requestAttachment.hasAuthentication()) {
+            HttpHandler httpHandler = RequestUtils.wrapSecurity(Application.getInstance(LocaleHandler.class), this.requestAttachment.getUsername(), this.requestAttachment.getPassword());
+            httpHandler.handleRequest(exchange);
+        } else {
+            Application.getInstance(LocaleHandler.class).handleRequest(exchange);    
+        }
+    }
+}