Add validation to check max int limit (#159)

Signed-off-by: John Mazanec <jmazane@amazon.com>
opensearch-project · Oct 27, 2021 · 769a137 · 769a137
1 parent 1b4b207
commit 769a137
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java
@@ -14,6 +14,7 @@
 import com.google.common.collect.ImmutableList;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.common.xcontent.XContentParser;
+import org.opensearch.index.mapper.NumberFieldMapper;
 import org.opensearch.knn.index.KNNMethodContext;
 import org.opensearch.knn.plugin.KNNPlugin;
 import org.opensearch.knn.plugin.transport.TrainingJobRouterAction;
@@ -105,11 +106,11 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr
             } else if (KNN_METHOD.equals(fieldName) && ensureNotSet(fieldName, knnMethodContext)) {
                 knnMethodContext = KNNMethodContext.parse(parser.map());
             } else if (DIMENSION.equals(fieldName) && ensureNotSet(fieldName, dimension)) {
-                dimension = parser.intValue();
+                dimension = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false);
             } else if (MAX_VECTOR_COUNT_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, maximumVectorCount)) {
-                maximumVectorCount = parser.intValue();
+                maximumVectorCount = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false);
             } else if (SEARCH_SIZE_PARAMETER.equals(fieldName) && ensureNotSet(fieldName, searchSize)) {
-                searchSize = parser.intValue();
+                searchSize = (Integer) NumberFieldMapper.NumberType.INTEGER.parse(parser.objectBytes(), false);
             } else if (MODEL_DESCRIPTION.equals(fieldName) && ensureNotSet(fieldName, description)) {
                 description = parser.text();
             } else {