AYS-379 | Rate Limit Feature Has Been Integrated (#380)

pom.xml

@@ -39,6 +39,9 @@
         <jsonwebtoken.version>0.12.6</jsonwebtoken.version>
         <bouncycastle.version>1.78.1</bouncycastle.version>
 
+        <bucket4j-core.version>8.10.1</bucket4j-core.version>
+        <guava.version>33.3.0-jre</guava.version>
+
         <objenesis.version>3.3</objenesis.version>
         <commons-compress.version>1.27.0</commons-compress.version>
         <commons-text.version>1.12.0</commons-text.version>
@@ -119,6 +122,25 @@
         <!-- ===================== -->
 
 
+        <!-- ======================= -->
+        <!-- rate limit dependencies -->
+        <!-- ======================= -->
+        <dependency>
+            <groupId>com.bucket4j</groupId>
+            <artifactId>bucket4j-core</artifactId>
+            <version>${bucket4j-core.version}</version>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+            <version>${guava.version}</version>
+        </dependency>
+        <!-- ======================= -->
+        <!-- rate limit dependencies -->
+        <!-- ======================= -->
+
+
         <!-- ===================== -->
         <!-- security dependencies -->
         <!-- ===================== -->

src/main/java/org/ays/auth/filter/AysBearerTokenAuthenticationFilter.java

@@ -1,21 +1,32 @@
 package org.ays.auth.filter;
 
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import io.github.bucket4j.Bandwidth;
+import io.github.bucket4j.Bucket;
 import jakarta.servlet.FilterChain;
 import jakarta.servlet.ServletException;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
+import jakarta.validation.constraints.NotNull;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.ays.auth.model.AysToken;
 import org.ays.auth.service.AysInvalidTokenService;
 import org.ays.auth.service.AysTokenService;
+import org.ays.common.util.HttpServletRequestUtil;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.http.HttpHeaders;
 import org.springframework.lang.NonNull;
 import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.stereotype.Component;
 import org.springframework.web.filter.OncePerRequestFilter;
 
 import java.io.IOException;
+import java.time.Duration;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 
 /**
  * AysBearerTokenAuthenticationFilter is a filter that intercepts HTTP requests and processes the Bearer tokens included in the Authorization headers.
@@ -32,13 +43,50 @@ public class AysBearerTokenAuthenticationFilter extends OncePerRequestFilter {
     private final AysInvalidTokenService invalidTokenService;
 
 
+    @Value("${ays.rate-limit.authorized.enabled}")
+    private boolean isAuthorizedRateLimitEnabled;
+
+    private static final int MAXIMUM_AUTHORIZED_REQUESTS_COUNTS = 20;
+    private static final int MAXIMUM_AUTHORIZED_REQUESTS_DURATION_MINUTES = 1;
+    private static final Duration MAXIMUM_AUTHORIZED_REQUESTS_DURATION = Duration.ofMinutes(MAXIMUM_AUTHORIZED_REQUESTS_DURATION_MINUTES);
+    private final LoadingCache<String, Bucket> authorizedBuckets = CacheBuilder.newBuilder()
+            .expireAfterWrite(MAXIMUM_AUTHORIZED_REQUESTS_DURATION_MINUTES, TimeUnit.MINUTES)
+            .build(new CacheLoader<>() {
+                @Override
+                public @NotNull Bucket load(@NotNull String key) {
+                    return newBucket(
+                            MAXIMUM_AUTHORIZED_REQUESTS_COUNTS,
+                            MAXIMUM_AUTHORIZED_REQUESTS_DURATION
+                    );
+                }
+            });
+
+
+    @Value("${ays.rate-limit.unauthorized.enabled}")
+    private boolean isUnauthorizedRateLimitEnabled;
+
+    private static final int MAXIMUM_UNAUTHORIZED_REQUESTS_COUNTS = 5;
+    private static final int MAXIMUM_UNAUTHORIZED_REQUESTS_DURATION_MINUTES = 10;
+    private static final Duration MAXIMUM_UNAUTHORIZED_REQUESTS_DURATION = Duration.ofMinutes(MAXIMUM_UNAUTHORIZED_REQUESTS_DURATION_MINUTES);
+    private final LoadingCache<String, Bucket> unauthorizedBuckets = CacheBuilder.newBuilder()
+            .expireAfterWrite(MAXIMUM_UNAUTHORIZED_REQUESTS_DURATION_MINUTES, TimeUnit.MINUTES)
+            .build(new CacheLoader<>() {
+                @Override
+                public @NotNull Bucket load(@NotNull String key) {
+                    return newBucket(
+                            MAXIMUM_UNAUTHORIZED_REQUESTS_COUNTS,
+                            MAXIMUM_UNAUTHORIZED_REQUESTS_DURATION
+                    );
+                }
+            });
+
+
     @Override
-    protected void doFilterInternal(HttpServletRequest httpServletRequest,
+    protected void doFilterInternal(@NotNull HttpServletRequest httpServletRequest,
                                     @NonNull HttpServletResponse httpServletResponse,
                                     @NonNull FilterChain filterChain) throws ServletException, IOException {
 
-
-        log.debug("API Request was secured with AYS Security!");
+        final String ipAddress = HttpServletRequestUtil.getClientIpAddress(httpServletRequest);
 
         final String authorizationHeader = httpServletRequest.getHeader(HttpHeaders.AUTHORIZATION);
         if (AysToken.isBearerToken(authorizationHeader)) {
@@ -49,11 +97,62 @@ protected void doFilterInternal(HttpServletRequest httpServletRequest,
             final String tokenId = tokenService.getPayload(jwt).getId();
             invalidTokenService.checkForInvalidityOfToken(tokenId);
 
+            if (isAuthorizedRateLimitEnabled) {
+                boolean isRateLimitExceeded = this.isRateLimitExceeded(ipAddress, authorizedBuckets, httpServletResponse);
+                if (isRateLimitExceeded) {
+                    return;
+                }
+            }
+
             final var authentication = tokenService.getAuthentication(jwt);
             SecurityContextHolder.getContext().setAuthentication(authentication);
+            filterChain.doFilter(httpServletRequest, httpServletResponse);
+            return;
+        }
+
+
+        if (isUnauthorizedRateLimitEnabled) {
+            boolean isRateLimitExceeded = this.isRateLimitExceeded(ipAddress, unauthorizedBuckets, httpServletResponse);
+            if (isRateLimitExceeded) {
+                return;
+            }
         }
 
         filterChain.doFilter(httpServletRequest, httpServletResponse);
+
+    }
+
+    private boolean isRateLimitExceeded(final String ipAddress,
+                                        final LoadingCache<String, Bucket> buckets,
+                                        final HttpServletResponse httpServletResponse) {
+
+        try {
+
+            final Bucket bucket = buckets.get(ipAddress);
+            if (bucket.tryConsume(1)) {
+                return false;
+            }
+
+        } catch (ExecutionException exception) {
+            log.error("Error while checking rate limit for IP: {}", ipAddress, exception);
+            httpServletResponse.setStatus(429);
+            return true;
+        }
+
+        log.warn("Rate limit exceeded for IP: {}", ipAddress);
+        httpServletResponse.setStatus(429);
+        return true;
+    }
+
+    private static Bucket newBucket(int maximumRequestsCounts, Duration maximumDuration) {
+        final Bandwidth bandwidth = Bandwidth
+                .builder()
+                .capacity(maximumRequestsCounts)
+                .refillIntervally(maximumRequestsCounts, maximumDuration)
+                .build();
+        return Bucket.builder()
+                .addLimit(bandwidth)
+                .build();
     }
 
 }

src/main/java/org/ays/common/util/HttpServletRequestUtil.java

@@ -0,0 +1,19 @@
+package org.ays.common.util;
+
+import jakarta.servlet.http.HttpServletRequest;
+import lombok.experimental.UtilityClass;
+
+@UtilityClass
+public class HttpServletRequestUtil {
+
+    public static String getClientIpAddress(final HttpServletRequest httpServletRequest) {
+
+        final String ipAddress = httpServletRequest.getHeader("X-Forwarded-For");
+        if (ipAddress == null || ipAddress.isEmpty()) {
+            return httpServletRequest.getRemoteAddr().trim();
+        }
+
+        return ipAddress.split(",")[0].trim();
+    }
+
+}

src/main/resources/application.yml

@@ -81,3 +81,8 @@ ays:
     invalid-tokens-deletion:
       cron: ${INVALID_TOKENS_DELETION_CRON:0 0 */3 * * ?}
       enable: ${INVALID_TOKENS_DELETION_ENABLED:true}
+  rate-limit:
+    authorized:
+      enabled: ${AYS_AUTHORIZED_RATE_LIMIT_ENABLED:false}
+    unauthorized:
+      enabled: ${AYS_UNAUTHORIZED_RATE_LIMIT_ENABLED:false}