Rebase, add TODOs

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Feb 23, 2024 · 8b2c8fe · 8b2c8fe
1 parent 6a169a8
commit 8b2c8fe
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/HttpHostStep.java b/src/main/java/org/opensearch/flowframework/workflow/HttpHostStep.java
@@ -40,9 +40,12 @@ public PlainActionFuture<WorkflowData> execute(
         String currentNodeId,
         WorkflowData currentNodeInputs,
         Map<String, WorkflowData> outputs,
-        Map<String, String> previousNodeInputs
+        Map<String, String> previousNodeInputs,
+        Map<String, String> params
     ) {
         Set<String> requiredKeys = Set.of(SCHEME_FIELD, HOSTNAME_FIELD, PORT_FIELD);
+        // TODO Possibly add credentials fields here
+        // See ML Commons MLConnectorInput class and its usage
         Set<String> optionalKeys = Collections.emptySet();
 
         try {
@@ -51,11 +54,12 @@ public PlainActionFuture<WorkflowData> execute(
                 optionalKeys,
                 currentNodeInputs,
                 outputs,
-                previousNodeInputs
+                previousNodeInputs,
+                params
             );
 
             String scheme = validScheme(inputs.get(SCHEME_FIELD));
-            String hostname = inputs.get(HOSTNAME_FIELD).toString();
+            String hostname = validHostName(inputs.get(HOSTNAME_FIELD));
             int port = validPort(inputs.get(PORT_FIELD));
 
             HttpHost httpHost = new HttpHost(scheme, hostname, port);
@@ -84,6 +88,14 @@ private String validScheme(Object o) {
         throw new FlowFrameworkException("http_host scheme must be http or https", RestStatus.BAD_REQUEST);
     }
 
+    private String validHostName(Object o) {
+        // TODO Add validation:
+        // Prevent use of localhost or private IP address ranges
+        // See ML Commons MLHttpClientFactory.java methods for examples
+        // Possibly consider an allowlist of addresses
+        return o.toString();
+    }
+
     private int validPort(Object o) {
         try {
             int port = Integer.parseInt(o.toString());

diff --git a/src/test/java/org/opensearch/flowframework/workflow/HttpHostStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/HttpHostStepTests.java
@@ -38,6 +38,7 @@ public void testHttpHost() throws InterruptedException, ExecutionException {
             inputData.getNodeId(),
             inputData,
             Collections.emptyMap(),
+            Collections.emptyMap(),
             Collections.emptyMap()
         );
 
@@ -62,6 +63,7 @@ public void testBadScheme() {
             badSchemeData.getNodeId(),
             badSchemeData,
             Collections.emptyMap(),
+            Collections.emptyMap(),
             Collections.emptyMap()
         );
 
@@ -84,6 +86,7 @@ public void testBadPort() {
             badPortData.getNodeId(),
             badPortData,
             Collections.emptyMap(),
+            Collections.emptyMap(),
             Collections.emptyMap()
         );
 
@@ -106,6 +109,7 @@ public void testNoParsePort() {
             noParsePortData.getNodeId(),
             noParsePortData,
             Collections.emptyMap(),
+            Collections.emptyMap(),
             Collections.emptyMap()
         );