feat: add an option to ignore host IP and subscription detail

https://gravitee.atlassian.net/browse/APIM-3620
gravitee-io · Feb 29, 2024 · ba5d3ee · ba5d3ee
1 parent 55b065c
commit ba5d3ee
Show file tree

Hide file tree

Showing 12 changed files with 91 additions and 1 deletion.
diff --git a/gravitee-policy-quota/src/main/java/io/gravitee/policy/quota/RateLimitKeyFactory.java b/gravitee-policy-quota/src/main/java/io/gravitee/policy/quota/RateLimitKeyFactory.java
@@ -48,6 +48,14 @@ public static String createRateLimitKey(ExecutionContext executionContext, Quota
 
         StringBuilder key = new StringBuilder();
 
+        if (configuration.isUseKeyOnly() && configuration.getKey() != null && !configuration.getKey().isEmpty()) {
+            return key
+                .append(executionContext.getTemplateEngine().getValue(configuration.getKey(), String.class))
+                .append(KEY_SEPARATOR)
+                .append(RATE_LIMIT_TYPE)
+                .toString();
+        }
+
         String plan = (String) executionContext.getAttribute(ExecutionContext.ATTR_PLAN);
         if (plan != null) {
             key

diff --git a/...policy-quota/src/main/java/io/gravitee/policy/quota/configuration/QuotaConfiguration.java b/...policy-quota/src/main/java/io/gravitee/policy/quota/configuration/QuotaConfiguration.java
@@ -25,8 +25,8 @@
  * @author David BRASSELY (david.brassely at graviteesource.com)
  * @author GraviteeSource Team
  */
-@Builder
 @Data
+@Builder(toBuilder = true)
 @AllArgsConstructor
 @NoArgsConstructor
 public class QuotaConfiguration {
@@ -40,4 +40,6 @@ public class QuotaConfiguration {
     private ChronoUnit periodTimeUnit;
 
     private String key;
+
+    private boolean useKeyOnly;
 }
diff --git a/gravitee-policy-quota/src/main/resources/schemas/schema-form.json b/gravitee-policy-quota/src/main/resources/schemas/schema-form.json
@@ -27,6 +27,12 @@
                         "expression-language": true
                     }
                 },
+                "useKeyOnly": {
+                    "type": "boolean",
+                    "default": false,
+                    "title": "Use key only",
+                    "description": "Only uses the custom key to identify the consumer, regardless of the subscription and plan."
+                },
                 "limit": {
                     "type": "integer",
                     "title": "Max requests (static)",

diff --git a/gravitee-policy-quota/src/test/java/io/gravitee/policy/quota/RateLimitKeyFactoryTest.java b/gravitee-policy-quota/src/test/java/io/gravitee/policy/quota/RateLimitKeyFactoryTest.java
@@ -110,4 +110,17 @@ void should_use_the_hashcoded_resolved_path_when_defined() {
         // Then
         Assertions.assertThat(key).isEqualTo("plansubscription:q:148210906");
     }
+
+    @Test
+    void should_use_only_the_key_defined_when_enabled() {
+        // Given
+        context.setAttribute(ExecutionContext.ATTR_PLAN, "plan");
+        context.setAttribute(ExecutionContext.ATTR_SUBSCRIPTION_ID, "subscription");
+
+        // When
+        String key = RateLimitKeyFactory.createRateLimitKey(context, QuotaConfiguration.builder().key("key").useKeyOnly(true).build());
+
+        // Then
+        Assertions.assertThat(key).isEqualTo("key:q");
+    }
 }
diff --git a/...itee-policy-ratelimit/src/main/java/io/gravitee/policy/ratelimit/RateLimitKeyFactory.java b/...itee-policy-ratelimit/src/main/java/io/gravitee/policy/ratelimit/RateLimitKeyFactory.java
@@ -48,6 +48,14 @@ public static String createRateLimitKey(ExecutionContext executionContext, RateL
 
         StringBuilder key = new StringBuilder();
 
+        if (configuration.isUseKeyOnly() && configuration.getKey() != null && !configuration.getKey().isEmpty()) {
+            return key
+                .append(executionContext.getTemplateEngine().getValue(configuration.getKey(), String.class))
+                .append(KEY_SEPARATOR)
+                .append(RATE_LIMIT_TYPE)
+                .toString();
+        }
+
         String plan = (String) executionContext.getAttribute(ExecutionContext.ATTR_PLAN);
         if (plan != null) {
             key

diff --git a/...imit/src/main/java/io/gravitee/policy/ratelimit/configuration/RateLimitConfiguration.java b/...imit/src/main/java/io/gravitee/policy/ratelimit/configuration/RateLimitConfiguration.java
@@ -40,4 +40,6 @@ public class RateLimitConfiguration {
     private TimeUnit periodTimeUnit;
 
     private String key;
+
+    private boolean useKeyOnly;
 }
diff --git a/gravitee-policy-ratelimit/src/main/resources/schemas/schema-form.json b/gravitee-policy-ratelimit/src/main/resources/schemas/schema-form.json
@@ -27,6 +27,12 @@
                         "expression-language": true
                     }
                 },
+                "useKeyOnly": {
+                    "type": "boolean",
+                    "default": false,
+                    "title": "Use key only",
+                    "description": "Only uses the custom key to identify the consumer, regardless of the subscription and plan."
+                },
                 "limit": {
                     "type": "integer",
                     "title": "Max requests (static)",

diff --git a/...-policy-ratelimit/src/test/java/io/gravitee/policy/ratelimit/RateLimitKeyFactoryTest.java b/...-policy-ratelimit/src/test/java/io/gravitee/policy/ratelimit/RateLimitKeyFactoryTest.java
@@ -110,4 +110,17 @@ void should_use_the_hashcoded_resolved_path_when_defined() {
         // Then
         Assertions.assertThat(key).isEqualTo("plansubscription:rl:148210906");
     }
+
+    @Test
+    void should_use_only_the_key_defined_when_enabled() {
+        // Given
+        context.setAttribute(ExecutionContext.ATTR_PLAN, "plan");
+        context.setAttribute(ExecutionContext.ATTR_SUBSCRIPTION_ID, "subscription");
+
+        // When
+        String key = RateLimitKeyFactory.createRateLimitKey(context, RateLimitConfiguration.builder().key("key").useKeyOnly(true).build());
+
+        // Then
+        Assertions.assertThat(key).isEqualTo("key:rl");
+    }
 }
diff --git a/gravitee-policy-spikearrest/src/main/java/io/gravitee/policy/spike/RateLimitKeyFactory.java b/gravitee-policy-spikearrest/src/main/java/io/gravitee/policy/spike/RateLimitKeyFactory.java
@@ -48,6 +48,14 @@ public static String createRateLimitKey(ExecutionContext executionContext, Spike
 
         StringBuilder key = new StringBuilder();
 
+        if (configuration.isUseKeyOnly() && configuration.getKey() != null && !configuration.getKey().isEmpty()) {
+            return key
+                .append(executionContext.getTemplateEngine().getValue(configuration.getKey(), String.class))
+                .append(KEY_SEPARATOR)
+                .append(RATE_LIMIT_TYPE)
+                .toString();
+        }
+
         String plan = (String) executionContext.getAttribute(ExecutionContext.ATTR_PLAN);
         if (plan != null) {
             key

diff --git a/...arrest/src/main/java/io/gravitee/policy/spike/configuration/SpikeArrestConfiguration.java b/...arrest/src/main/java/io/gravitee/policy/spike/configuration/SpikeArrestConfiguration.java
@@ -40,4 +40,6 @@ public class SpikeArrestConfiguration {
     private TimeUnit periodTimeUnit;
 
     private String key;
+
+    private boolean useKeyOnly;
 }
diff --git a/gravitee-policy-spikearrest/src/main/resources/schemas/schema-form.json b/gravitee-policy-spikearrest/src/main/resources/schemas/schema-form.json
@@ -27,6 +27,12 @@
                         "expression-language": true
                     }
                 },
+                "useKeyOnly": {
+                    "type": "boolean",
+                    "default": false,
+                    "title": "Use key only",
+                    "description": "Only uses the custom key to identify the consumer, regardless of the subscription and plan."
+                },
                 "limit": {
                     "type": "integer",
                     "title": "Max requests (static)",

diff --git a/...ee-policy-spikearrest/src/test/java/io/gravitee/policy/spike/RateLimitKeyFactoryTest.java b/...ee-policy-spikearrest/src/test/java/io/gravitee/policy/spike/RateLimitKeyFactoryTest.java
@@ -110,4 +110,20 @@ void should_use_the_hashcoded_resolved_path_when_defined() {
         // Then
         Assertions.assertThat(key).isEqualTo("plansubscription:sa:148210906");
     }
+
+    @Test
+    void should_use_only_the_key_defined_when_enabled() {
+        // Given
+        context.setAttribute(ExecutionContext.ATTR_PLAN, "plan");
+        context.setAttribute(ExecutionContext.ATTR_SUBSCRIPTION_ID, "subscription");
+
+        // When
+        String key = RateLimitKeyFactory.createRateLimitKey(
+            context,
+            SpikeArrestConfiguration.builder().key("key").useKeyOnly(true).build()
+        );
+
+        // Then
+        Assertions.assertThat(key).isEqualTo("key:sa");
+    }
 }