From a9397553bb07def2dec27349fc259547e7d7add2 Mon Sep 17 00:00:00 2001
From: Naveen Tatikonda <navtat@amazon.com>
Date: Mon, 4 Apr 2022 18:07:34 -0500
Subject: [PATCH] Add size validation for Search Model API

Signed-off-by: Naveen Tatikonda <navtat@amazon.com>
---
 .../knn/plugin/rest/RestSearchModelHandler.java    | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/src/main/java/org/opensearch/knn/plugin/rest/RestSearchModelHandler.java b/src/main/java/org/opensearch/knn/plugin/rest/RestSearchModelHandler.java
index eeb510b9d3..4e76158325 100644
--- a/src/main/java/org/opensearch/knn/plugin/rest/RestSearchModelHandler.java
+++ b/src/main/java/org/opensearch/knn/plugin/rest/RestSearchModelHandler.java
@@ -14,6 +14,7 @@
 import com.google.common.collect.ImmutableList;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.client.node.NodeClient;
+import org.opensearch.common.ValidationException;
 import org.opensearch.knn.plugin.KNNPlugin;
 import org.opensearch.knn.plugin.transport.SearchModelAction;
 import org.opensearch.rest.BaseRestHandler;
@@ -69,9 +70,22 @@ private void checkUnSupportedParamsExists(RestRequest request) {
         throw new IllegalArgumentException(errorMessage);
     }
 
+    private void searchSizeValidation(RestRequest request) {
+        ValidationException validationException = new ValidationException();
+        if (request.hasParam("size")) {
+            int searchSize = request.paramAsInt("size", 1);
+            if (searchSize > 1000 || searchSize <= 0) {
+                String errorMessage = "Size must be > 0 and <= 1000";
+                validationException.addValidationError(errorMessage);
+                throw validationException;
+            }
+        }
+    }
+
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         checkUnSupportedParamsExists(request);
+        searchSizeValidation(request);
         SearchRequest searchRequest = new SearchRequest();
         IntConsumer setSize = size -> searchRequest.source().size(size);
         request.withContentOrSourceParamParserOrNull(