Customize rate limit settings in kustvakt.conf

[AI assisted]

Change-Id: I5907dbd839151b00bab2dfea9dedc3d5ffb11765

Author: margaretha

Changes

@@ -6,7 +6,8 @@
 - Change Userdata to use String username instead of integer userId
 - Allow admin to create groups with name length less than 3 characters 
   to support existing groups from C2
-- Implemented rate limit for authenticated users (with AI assistance)
+- Implemented rate limit for authenticated users (AI assisted)
+- Customize rate limit settings in kustvakt.conf (AI assisted)
 
 # version 1.1
 

src/main/java/de/ids_mannheim/korap/web/filter/RateLimitFilter.java

@@ -5,27 +5,34 @@
 import java.time.Duration;
 import java.util.Base64;
 import java.util.Objects;
+import java.util.Properties;
 import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 
+import de.ids_mannheim.korap.config.KustvaktConfiguration;
+import de.ids_mannheim.korap.utils.TimeUtils;
+import jakarta.annotation.PostConstruct;
 import jakarta.annotation.Priority;
 import jakarta.ws.rs.Priorities;
 import jakarta.ws.rs.WebApplicationException;
 import jakarta.ws.rs.container.ContainerRequestContext;
 import jakarta.ws.rs.container.ContainerRequestFilter;
 import jakarta.ws.rs.core.HttpHeaders;
 import jakarta.ws.rs.core.Response;
+import lombok.Getter;
 
-/** Implemented with AI assistance
- * 
+/**
  * Simple in-memory rate limitation for authenticated users.
  * <p>
  * Keyed by bearer token (preferred) or username (fallback).
  * <p>
  * Note: In-memory means per-JVM only. For clustered deployments, use Redis/etc.
+ * 
+ * Implemented with AI assistance
  */
 @Component
 @Priority(Priorities.AUTHORIZATION)
@@ -34,53 +41,96 @@ public class RateLimitFilter implements ContainerRequestFilter {
     private static final Logger jlog = LogManager
             .getLogger(RateLimitFilter.class);
 
-    // Defaults: 60 requests per minute per key
-    // Keep these conservative and easy to change later via config injection.
-    private static final long REFILL_TOKENS = 60;
-    private static final Duration REFILL_PERIOD = Duration.ofMinutes(1);
-    public static final long BURST_CAPACITY = 60;
-
-    /**
-     * Prevent unbounded growth: keep at most this many distinct keys in-memory.
-     */
-    private static final int MAX_BUCKETS = 10_000;
+    @Autowired
+    private KustvaktConfiguration config;
 
-    /**
-     * Evict buckets that haven't been seen for this long.
-     */
-    private static final Duration BUCKET_TTL = Duration.ofHours(6);
+    // Rate limiting configuration (loaded from kustvakt.conf)
+    private long refillTokens;
+    private Duration refillPeriod;
+    @Getter
+    private long burstCapacity;
+    private int maxBuckets;
+    private Duration bucketTTL;
 
     private final ConcurrentHashMap<String, BucketEntry> buckets = new ConcurrentHashMap<>();
 
+    @PostConstruct
+    private void initializeConfiguration() {
+        try {
+            Properties props = config.getProperties();
+            
+            // Handle case where properties might not be initialized yet
+            if (props == null) {
+                jlog.warn("KustvaktConfiguration properties not available, using default rate limiting values");
+                setDefaultValues();
+                return;
+            }
+            
+            // Load rate limiting settings from kustvakt.conf with sensible defaults
+            String refillTokensStr = props.getProperty("ratelimit.refill.tokens", "60");
+            this.refillTokens = Long.parseLong(refillTokensStr);
+
+            String refillPeriodStr = props.getProperty("ratelimit.refill.period", "1M");
+            this.refillPeriod = Duration.ofSeconds(TimeUtils.convertTimeToSeconds(refillPeriodStr));
+
+            String burstCapacityStr = props.getProperty("ratelimit.burst.capacity", "60");
+            this.burstCapacity = Long.parseLong(burstCapacityStr);
+
+            String maxBucketsStr = props.getProperty("ratelimit.max.buckets", "10000");
+            this.maxBuckets = Integer.parseInt(maxBucketsStr);
+
+            String bucketTTLStr = props.getProperty("ratelimit.bucket.ttl", "6H");
+            this.bucketTTL = Duration.ofSeconds(TimeUtils.convertTimeToSeconds(bucketTTLStr));
+
+            jlog.info("Rate limiting initialized: refillTokens={}, refillPeriod={}, burstCapacity={}, maxBuckets={}, bucketTTL={}",
+                    refillTokens, refillPeriod, burstCapacity, maxBuckets, bucketTTL);
+        } catch (Exception e) {
+            jlog.error("Failed to initialize rate limiting configuration, using defaults", e);
+            setDefaultValues();
+        }
+    }
+    
+    private void setDefaultValues() {
+        this.refillTokens = 60;
+        this.refillPeriod = Duration.ofMinutes(1);
+        this.burstCapacity = 60;
+        this.maxBuckets = 10000;
+        this.bucketTTL = Duration.ofHours(6);
+        jlog.info("Rate limiting initialized with defaults: refillTokens={}, refillPeriod={}, burstCapacity={}, maxBuckets={}, bucketTTL={}",
+                refillTokens, refillPeriod, burstCapacity, maxBuckets, bucketTTL);
+    }
+
     @Override
     public void filter (ContainerRequestContext request) {
         // Only apply to authenticated requests
         if (request.getSecurityContext() == null
                 || request.getSecurityContext().getUserPrincipal() == null) {
+            jlog.debug("Skipping rate limiting - no SecurityContext or UserPrincipal");
             return;
         }
 
         String key = resolveKey(request);
         if (key == null) {
+            jlog.debug("Skipping rate limiting - could not resolve key");
             return;
         }
 
+        jlog.debug("Applying rate limiting for key: {}", key);
+
         long now = System.currentTimeMillis();
 
         // Opportunistic cleanup to avoid memory growth.
-        // Do it only on inserts or if we grow too large.
-        if (buckets.size() > MAX_BUCKETS) {
+        if (buckets.size() > maxBuckets) {
             cleanupOldEntries(now);
         }
 
         BucketEntry entry = buckets.compute(key, (k, existing) -> {
             if (existing == null) {
-                // If we're still too large, try another cleanup pass before adding.
-                if (buckets.size() > MAX_BUCKETS) {
+                if (buckets.size() > maxBuckets) {
                     cleanupOldEntries(now);
                 }
-                return new BucketEntry(new TokenBucket(BURST_CAPACITY,
-                        REFILL_TOKENS, REFILL_PERIOD.toMillis()), now);
+                return new BucketEntry(new TokenBucket(burstCapacity,
+                        refillTokens, refillPeriod.toMillis()), now);
             }
             existing.lastSeenAtMillis = now;
             return existing;
@@ -90,20 +140,29 @@ public void filter (ContainerRequestContext request) {
             long retryAfterSeconds = Math
                     .max(1, entry.bucket.millisUntilNextToken() / 1000);
 
+            jlog.info("Rate limit exceeded for key: {}, retry after {} seconds", key, retryAfterSeconds);
             throw new WebApplicationException(Response.status(429)
                     .header("Retry-After", String.valueOf(retryAfterSeconds))
                     .entity("Rate limit exceeded")
                     .build());
         }
     }
 
+    /**
+     * Clear all rate limit buckets. For testing purposes only.
+     */
+    public void clearBuckets() {
+        buckets.clear();
+        jlog.info("Rate limit buckets cleared");
+    }
+
     private void cleanupOldEntries (long nowMillis) {
-        final long cutoff = nowMillis - BUCKET_TTL.toMillis();
+        final long cutoff = nowMillis - bucketTTL.toMillis();
         buckets.entrySet().removeIf(e -> e.getValue().lastSeenAtMillis < cutoff);
 
         // Still too big? Remove arbitrary entries (best-effort bound).
-        if (buckets.size() > MAX_BUCKETS) {
-            int toRemove = buckets.size() - MAX_BUCKETS;
+        if (buckets.size() > maxBuckets) {
+            int toRemove = buckets.size() - maxBuckets;
             for (String k : buckets.keySet()) {
                 buckets.remove(k);
                 if (--toRemove <= 0)
@@ -123,6 +182,8 @@ private String resolveKey (ContainerRequestContext request) {
             }
         }
 
+        // Skip unauthenticated requests. DemoUserFilter sets username guest 
+        // for such requests.
         // Fallback to username/principal name
 //        String name = request.getSecurityContext().getUserPrincipal().getName();
 //        if (name != null && !name.isBlank()) {
@@ -208,4 +269,4 @@ private static final class BucketEntry {
             this.lastSeenAtMillis = lastSeenAtMillis;
         }
     }
-}
+}

src/main/resources/kustvakt.conf

@@ -59,6 +59,19 @@ availability.regex.free = CC.*
 availability.regex.public = ACA.*|QAO-NC
 availability.regex.all = QAO.*
 
+# Rate limiting for authenticated users
+#
+# Number of requests allowed per time period
+ratelimit.refill.tokens = 60
+# Time period for token refill (format: 1S, 30M, 1H, 1D)
+ratelimit.refill.period = 1M
+# Maximum burst capacity (tokens that can be consumed immediately)
+ratelimit.burst.capacity = 60
+# Maximum number of rate limit buckets to keep in memory
+ratelimit.max.buckets = 10000
+# Time to live for unused rate limit buckets
+ratelimit.bucket.ttl = 6H
+
 # options referring to the security module!
 
 # OAuth 

src/test/java/de/ids_mannheim/korap/web/controller/RateLimitTest.java

@@ -2,7 +2,9 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
 
 import com.fasterxml.jackson.databind.JsonNode;
 
@@ -13,12 +15,23 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.Response.Status;
 
-/**
+/**  
  * Verifies authenticated rate limiting (HTTP 429) is applied after
  * auth.
+ * 
+ * Implemented with AI assistance
  */
 public class RateLimitTest extends OAuth2TestBase {
+	@Autowired
+    private RateLimitFilter rateLimitFilter;
 
+    @BeforeEach
+    public void clearRateLimitState() {
+        // Clear rate limit state before each test
+        if (rateLimitFilter != null) {
+            rateLimitFilter.clearBuckets();
+        }
+    }
 	@Test
 	public void testAuthenticatedRateLimitBearerToken ()
 			throws KustvaktException {
@@ -27,7 +40,7 @@ public void testAuthenticatedRateLimitBearerToken ()
 		JsonNode node = JsonUtils.readTree(response.readEntity(String.class));
 		String accessToken = node.at("/access_token").asText();
 
-		for (long i = 0; i < RateLimitFilter.BURST_CAPACITY; i++) {
+		for (long i = 0; i < rateLimitFilter.getBurstCapacity(); i++) {
 			Response r = searchWithAccessToken(accessToken);
 			assertEquals(Status.OK.getStatusCode(), r.getStatus(),
 					"request " + i);

src/test/java/de/ids_mannheim/korap/web/lite/RateLimitAnonymousTest.java

@@ -8,8 +8,10 @@
 import jakarta.ws.rs.core.Response;
 import jakarta.ws.rs.core.Response.Status;
 
-/**
+/** 
  * Verifies unauthenticated requests are not rate-limited.
+ * 
+ * Implemented with AI assistance
  */
 public class RateLimitAnonymousTest extends LiteJerseyTest {
 

src/test/resources/kustvakt-test.conf

@@ -61,6 +61,18 @@ availability.regex.public = ACA.*|QAO-NC
 # availability.regex.all = QAO.*
 availability.regex.all = QAO-NC-LOC:ids.*
 
+# Rate limiting for authenticated users
+#
+# Number of requests allowed per time period
+ratelimit.refill.tokens = 5
+# Time period for token refill (format: 1S, 30M, 1H, 1D)
+ratelimit.refill.period = 1M
+# Maximum burst capacity (tokens that can be consumed immediately)
+ratelimit.burst.capacity = 5
+# Maximum number of rate limit buckets to keep in memory
+ratelimit.max.buckets = 10000
+# Time to live for unused rate limit buckets
+ratelimit.bucket.ttl = 6H
 
 # options referring to the security module!