Add HTTP requests rate limiter

Signed-off-by: Paolo Di Tommaso <paolo.ditommaso@gmail.com>
seqeralabs · Dec 15, 2024 · 632ba33 · 632ba33
1 parent b2e31bf
commit 632ba33
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 5 deletions.
diff --git a/src/main/groovy/io/seqera/wave/configuration/RateLimiterConfig.groovy b/src/main/groovy/io/seqera/wave/configuration/RateLimiterConfig.groovy
@@ -71,4 +71,17 @@ interface RateLimiterConfig {
 
     }
 
+    RequestsLimit getRequest()
+
+    @ConfigurationProperties('requests')
+    static interface RequestsLimit {
+
+        @Bindable("10 / 1")
+        LimitConfig getAnonymous()
+
+        @Bindable("100 / 1s")
+        LimitConfig getAuthenticated()
+
+    }
+
 }
diff --git a/src/main/groovy/io/seqera/wave/controller/RegistryProxyController.groovy b/src/main/groovy/io/seqera/wave/controller/RegistryProxyController.groovy
@@ -121,9 +121,14 @@ class RegistryProxyController {
         log.info "> Request [$httpRequest.method] $httpRequest.path"
         final route = routeHelper.parse("/v2/" + url)
 
-        if( route.manifest && route.digest ){
-            String ip = addressResolver.resolve(httpRequest)
-            rateLimiterService?.acquirePull( new AcquireRequest(route.identity.userEmail, ip) )
+        if( rateLimiterService ) {
+            final ip = addressResolver.resolve(httpRequest)
+            final id = new AcquireRequest(route.identity.userEmail, ip)
+            rateLimiterService.acquireRequest(id)
+
+            if( route.manifest && route.digest ){
+                rateLimiterService.acquirePull(id)
+            }
         }
 
         // check if it's a container under build

diff --git a/src/main/groovy/io/seqera/wave/ratelimit/RateLimiterService.groovy b/src/main/groovy/io/seqera/wave/ratelimit/RateLimiterService.groovy
@@ -19,8 +19,6 @@
 package io.seqera.wave.ratelimit
 
 import io.seqera.wave.exception.SlowDownException
-
-
 /**
  * @author : jorge <jorge.aguilera@seqera.io>
  *
@@ -31,5 +29,7 @@ interface RateLimiterService {
 
     void acquirePull(AcquireRequest request) throws SlowDownException
 
+    void acquireRequest(AcquireRequest request) throws SlowDownException
+
     boolean acquireTimeoutCounter(String endpoint)
 }
diff --git a/src/main/groovy/io/seqera/wave/ratelimit/impl/SpillwayRateLimiter.groovy b/src/main/groovy/io/seqera/wave/ratelimit/impl/SpillwayRateLimiter.groovy
@@ -31,6 +31,7 @@ import io.seqera.wave.configuration.RateLimiterConfig
 import io.seqera.wave.exception.SlowDownException
 import io.seqera.wave.ratelimit.AcquireRequest
 import io.seqera.wave.ratelimit.RateLimiterService
+import io.seqera.wave.tower.auth.JwtAuth
 import jakarta.inject.Singleton
 import jakarta.validation.constraints.NotNull
 /**
@@ -55,6 +56,11 @@ class SpillwayRateLimiter implements RateLimiterService {
 
     Spillway<String> timeoutErrors
 
+    Spillway<String> anonymousRequests
+
+    Spillway<String> authsRequests
+
+
     SpillwayRateLimiter(@NotNull LimitUsageStorage storage, @NotNull RateLimiterConfig config) {
         init(storage, config)
     }
@@ -63,6 +69,7 @@ class SpillwayRateLimiter implements RateLimiterService {
         SpillwayFactory spillwayFactory = new SpillwayFactory(storage)
         initBuilds(spillwayFactory, config)
         initPulls(spillwayFactory, config)
+        initRequests(spillwayFactory, config)
         initTimeoutErrors(spillwayFactory, config)
     }
 
@@ -86,6 +93,16 @@ class SpillwayRateLimiter implements RateLimiterService {
         }
     }
 
+    @Override
+    void acquireRequest(AcquireRequest request) throws SlowDownException {
+        Spillway<String> resource = request.user ? authsRequests : anonymousRequests
+        String key = request.user ?: request.ip
+        if (!resource.tryCall(key)) {
+            final prefix = request.user ? 'user' : 'IP'
+            throw new SlowDownException("Request exceeded HTTP rate limit for $prefix $key")
+        }
+    }
+
     @Override
     boolean acquireTimeoutCounter(String endpoint) {
         try {
@@ -120,6 +137,14 @@ class SpillwayRateLimiter implements RateLimiterService {
         timeoutErrors = createLimit("timeoutErrors", spillwayFactory, config.timeoutErrors.getMaxRate())
     }
 
+    private void initRequests(SpillwayFactory spillwayFactory, RateLimiterConfig config) {
+        log.info "Http requests anonymous rate limit: max=$config.request.anonymous.max; duration:$config.request.anonymous.duration"
+        anonymousRequests = createLimit("anonymousRequests", spillwayFactory, config.request.anonymous)
+
+        log.info "Builds auth rate limit: max=$config.request.authenticated.max; duration:$config.request.authenticated.duration"
+        authsRequests = createLimit("authenticatedRequests", spillwayFactory, config.request.authenticated)
+    }
+
     private static Spillway<String> createLimit(String name, SpillwayFactory spillwayFactory, LimitConfig config) {
         Limit<String> userLimit = LimitBuilder.of("perUser")
                 .to(config.max)