Add verification to rate limiter number field (opensearch-project#2113)

* Add verification to rate limiter number field Signed-off-by: Sicheng Song <sicheng.song@outlook.com> * Change int to double Signed-off-by: Sicheng Song <sicheng.song@outlook.com> --------- Signed-off-by: Sicheng Song <sicheng.song@outlook.com>
zane-neo · Feb 20, 2024 · 0903d5d · 0903d5d
1 parent 7bcef02
commit 0903d5d
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/controller/MLRateLimiter.java b/common/src/main/java/org/opensearch/ml/common/controller/MLRateLimiter.java
@@ -8,6 +8,7 @@
 import lombok.Builder;
 import lombok.Getter;
 import lombok.Setter;
+import org.opensearch.OpenSearchParseException;
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
@@ -49,6 +50,14 @@ public static MLRateLimiter parse(XContentParser parser) throws IOException {
             switch (fieldName) {
                 case LIMIT_FIELD:
                     limit = parser.text();
+                    try {
+                        double limitNumber = Double.parseDouble(limit);
+                        if (limitNumber < 0) {
+                            throw new OpenSearchParseException("Limit field must be a positive number.");
+                        }
+                    } catch (NumberFormatException e) {
+                        throw new OpenSearchParseException("Limit field must be a positive number.");
+                    }
                     break;
                 case UNIT_FIELD:
                     unit = TimeUnit.valueOf(parser.text());

diff --git a/common/src/test/java/org/opensearch/ml/common/controller/MLRateLimiterTest.java b/common/src/test/java/org/opensearch/ml/common/controller/MLRateLimiterTest.java
@@ -20,6 +20,7 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
+import org.opensearch.OpenSearchParseException;
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.common.settings.Settings;
@@ -53,6 +54,7 @@ public void setUp() throws Exception {
                 .limit("1")
                 .unit(TimeUnit.MILLISECONDS)
                 .build();
+
         rateLimiterWithNumber = MLRateLimiter.builder()
                 .limit("1")
                 .build();
@@ -102,6 +104,20 @@ public void parseSuccess() throws Exception {
         });
     }
 
+    @Test
+    public void parseWithIllegalLimit() throws Exception {
+        exceptionRule.expect(OpenSearchParseException.class);
+        String inputStrWithIllegalLimit = "{\"limit\":\"-1\",\"unit\":\"MILLISECONDS\"}";
+        testParseFromJsonString(inputStrWithIllegalLimit, parsedInput -> {});
+    }
+
+    @Test
+    public void parseWithNegativeLimit() throws Exception {
+        exceptionRule.expect(OpenSearchParseException.class);
+        String inputStrWithNegativeLimit = "{\"limit\":\"ILLEGAL\",\"unit\":\"MILLISECONDS\"}";
+        testParseFromJsonString(inputStrWithNegativeLimit, parsedInput -> {});
+    }
+
     @Test
     public void parseWithNullField() throws Exception {
         exceptionRule.expect(IllegalStateException.class);