Fix connector private IP validation when executing agent without remote model (#3862)

* Fix bug in execute with private IP

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>

* Use isConnectorPrivateIpEnabled() from mlFeatureEnabledSetting

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>

* Remove redundant changes

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>

---------

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>

Author: nathaliellenaa

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java

@@ -90,6 +90,8 @@ public void invokeRemoteService_invalidIpAddress() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(false);
+        executor.setConnectorPrivateIpEnabled(privateIpEnabled);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -134,8 +136,27 @@ public void invokeRemoteService_EnabledPrivateIpAddress() {
                 actionListener
             );
         Mockito.verify(actionListener, never()).onFailure(any());
+    }
 
-        privateIpEnabled.set(false);
+    @Test
+    public void invokeRemoteService_DisabledPrivateIpAddress() {
+        ConnectorAction predictAction = ConnectorAction
+            .builder()
+            .actionType(PREDICT)
+            .method("POST")
+            .url("http://127.0.0.1/mock")
+            .requestBody("{\"input\": \"${parameters.input}\"}")
+            .build();
+        Connector connector = HttpConnector
+            .builder()
+            .name("test connector")
+            .version("1")
+            .protocol("http")
+            .actions(Arrays.asList(predictAction))
+            .build();
+        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(false);
+        executor.setConnectorPrivateIpEnabled(privateIpEnabled);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -168,6 +189,8 @@ public void invokeRemoteService_Empty_payload() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        AtomicBoolean privateIpEnabled = new AtomicBoolean(false);
+        executor.setConnectorPrivateIpEnabled(privateIpEnabled);
         executor.invokeRemoteService(PREDICT.name(), createMLInput(), new HashMap<>(), null, new ExecutionContext(0), actionListener);
         ArgumentCaptor<Exception> captor = ArgumentCaptor.forClass(IllegalArgumentException.class);
         Mockito.verify(actionListener, times(1)).onFailure(captor.capture());

plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java

@@ -18,6 +18,7 @@
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.connector.Connector;
 import org.opensearch.ml.common.connector.ConnectorAction;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.common.transport.connector.MLConnectorDeleteRequest;
 import org.opensearch.ml.common.transport.connector.MLExecuteConnectorAction;
@@ -43,6 +44,7 @@ public class ExecuteConnectorTransportAction extends HandledTransportAction<Acti
 
     ConnectorAccessControlHelper connectorAccessControlHelper;
     EncryptorImpl encryptor;
+    MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     @Inject
     public ExecuteConnectorTransportAction(
@@ -53,7 +55,8 @@ public ExecuteConnectorTransportAction(
         ScriptService scriptService,
         NamedXContentRegistry xContentRegistry,
         ConnectorAccessControlHelper connectorAccessControlHelper,
-        EncryptorImpl encryptor
+        EncryptorImpl encryptor,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
     ) {
         super(MLExecuteConnectorAction.NAME, transportService, actionFilters, MLConnectorDeleteRequest::new);
         this.client = client;
@@ -62,6 +65,7 @@ public ExecuteConnectorTransportAction(
         this.xContentRegistry = xContentRegistry;
         this.connectorAccessControlHelper = connectorAccessControlHelper;
         this.encryptor = encryptor;
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
     }
 
     @Override
@@ -77,6 +81,7 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
                     connector.decrypt(connectorAction, (credential, tenantId) -> encryptor.decrypt(credential, null), null);
                     RemoteConnectorExecutor connectorExecutor = MLEngineClassLoader
                         .initInstance(connector.getProtocol(), connector, Connector.class);
+                    connectorExecutor.setConnectorPrivateIpEnabled(mlFeatureEnabledSetting.isConnectorPrivateIpEnabled());
                     connectorExecutor.setScriptService(scriptService);
                     connectorExecutor.setClusterService(clusterService);
                     connectorExecutor.setClient(client);

plugin/src/test/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportActionTests.java

@@ -30,6 +30,7 @@
 import org.opensearch.ml.common.connector.Connector;
 import org.opensearch.ml.common.connector.ConnectorProtocols;
 import org.opensearch.ml.common.connector.HttpConnector;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.common.transport.connector.MLExecuteConnectorRequest;
 import org.opensearch.ml.engine.encryptor.EncryptorImpl;
@@ -70,6 +71,8 @@ public class ExecuteConnectorTransportActionTests extends OpenSearchTestCase {
     @Mock
     private EncryptorImpl encryptor;
     @Mock
+    private MLFeatureEnabledSetting mlFeatureEnabledSetting;
+    @Mock
     private HttpConnector connector;
     @Mock
     private Task task;
@@ -109,7 +112,8 @@ public void setup() {
             scriptService,
             xContentRegistry,
             connectorAccessControlHelper,
-            encryptor
+            encryptor,
+            mlFeatureEnabledSetting
         );
     }