Add tool security-related tests

Signed-off-by: Nathalie Jonathan <[email protected]>

Author: nathaliellenaa

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/tool/MLToolExecutorTest.java

@@ -202,4 +202,26 @@ public void test_ImmutableEmptyParametersMap() {
         Output output = outputCaptor.getValue();
         Assert.assertTrue(output instanceof ModelTensorOutput);
     }
+
+    @Test
+    public void test_ToolExecutionFailsWithoutProperPermission() {
+        when(toolMLInput.getToolName()).thenReturn("TestTool");
+        when(toolMLInput.getInputDataset()).thenReturn(inputDataSet);
+        when(inputDataSet.getParameters()).thenReturn(parameters);
+        when(toolFactory.create(any())).thenReturn(tool);
+        when(tool.validate(parameters)).thenReturn(true);
+
+        Mockito.doAnswer(invocation -> {
+            ActionListener<Object> listener = invocation.getArgument(1);
+            listener.onFailure(new SecurityException("no permissions for [indices:data/read/search] and User [name=test_user]"));
+            return null;
+        }).when(tool).run(Mockito.eq(parameters), any());
+
+        mlToolExecutor.execute(toolMLInput, actionListener);
+
+        Mockito.verify(actionListener).onFailure(exceptionCaptor.capture());
+        Exception exception = exceptionCaptor.getValue();
+        Assert.assertTrue(exception instanceof SecurityException);
+        Assert.assertTrue(exception.getMessage().contains("no permissions"));
+    }
 }

plugin/build.gradle

@@ -237,6 +237,7 @@ integTest {
         filter {
             excludeTestsMatching "org.opensearch.ml.rest.SecureMLRestIT"
             excludeTestsMatching "org.opensearch.ml.rest.MLModelGroupRestIT"
+            excludeTestsMatching "org.opensearch.ml.tools.ListIndexToolIT"
         }
     }
 

plugin/src/test/java/org/opensearch/ml/tools/ListIndexToolIT.java

@@ -14,9 +14,13 @@
 import java.util.Objects;
 
 import org.apache.commons.lang3.StringUtils;
+import org.apache.hc.core5.http.HttpHost;
 import org.junit.Before;
 import org.opensearch.client.Response;
+import org.opensearch.client.ResponseException;
+import org.opensearch.client.RestClient;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.commons.rest.SecureRestClientBuilder;
 import org.opensearch.ml.engine.tools.ListIndexTool;
 import org.opensearch.ml.rest.RestBaseAgentToolsIT;
 import org.opensearch.ml.utils.TestHelper;
@@ -37,6 +41,44 @@ public void setUpCluster() throws Exception {
         registerListIndexFlowAgent();
     }
 
+    public void testListIndexWithNoPermissions() throws Exception {
+        if (!isHttps()) {
+            log.info("Skipping permission test as security is not enabled");
+            return;
+        }
+
+        String noPermissionUser = "no_permission_user";
+        String password = "TestPassword123!";
+
+        try {
+            createUser(noPermissionUser, password, new ArrayList<>());
+
+            final RestClient noPermissionClient = new SecureRestClientBuilder(
+                getClusterHosts().toArray(new HttpHost[0]),
+                isHttps(),
+                noPermissionUser,
+                password
+            ).setSocketTimeout(60000).build();
+
+            try {
+                ResponseException exception = expectThrows(ResponseException.class, () -> {
+                    TestHelper
+                        .makeRequest(noPermissionClient, "POST", "/_plugins/_ml/agents/" + agentId + "/_execute", null, question, null);
+                });
+
+                String errorMessage = exception.getMessage().toLowerCase();
+                assertTrue(
+                    "Expected permission error, got: " + errorMessage,
+                    errorMessage.contains("no permissions") || errorMessage.contains("forbidden") || errorMessage.contains("unauthorized")
+                );
+            } finally {
+                noPermissionClient.close();
+            }
+        } finally {
+            deleteUser(noPermissionUser);
+        }
+    }
+
     private List<String> createIndices(int count) throws IOException {
         List<String> indices = new ArrayList<>();
         for (int i = 0; i < count; i++) {

plugin/src/test/java/org/opensearch/ml/tools/ScratchPadToolIT.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.tools;
+
+import static org.opensearch.ml.utils.TestHelper.makeRequest;
+
+import org.junit.Before;
+import org.opensearch.ml.rest.MLCommonsRestTestCase;
+import org.opensearch.test.OpenSearchIntegTestCase;
+
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 3)
+public class ScratchPadToolIT extends MLCommonsRestTestCase {
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+    }
+
+    public void testScratchpadSizeLimit() throws Exception {
+        String largeContent = "A".repeat(100 * 1024 * 1024);
+        String requestBody = String.format("{\"parameters\":{\"notes\":\"%s\"}}", largeContent);
+
+        Exception exception = expectThrows(Exception.class, () -> {
+            makeRequest(client(), "POST", "/_plugins/_ml/tools/_execute/WriteToScratchPadTool", null, requestBody, null);
+        });
+
+        String errorMessage = exception.getMessage().toLowerCase();
+        assertTrue(
+            "Expected HTTP content length error, got: " + errorMessage,
+            errorMessage.contains("content length")
+                || errorMessage.contains("too large")
+                || errorMessage.contains("entity too large")
+                || errorMessage.contains("413")
+        );
+    }
+}