add security IT (#168)

Signed-off-by: Yaliang Wu <[email protected]>

Author: ylwu-amzn

plugin/build.gradle

@@ -104,6 +104,12 @@ integTest {
         }
     }
 
+    if (System.getProperty("https") == null || System.getProperty("https") == "false") {
+        filter {
+            excludeTestsMatching "org.opensearch.ml.rest.SecureMLRestIT"
+        }
+    }
+
     // The 'doFirst' delays till execution time.
     doFirst {
         // Tell the test JVM if the cluster JVM is running under a debugger so that tests can
@@ -182,7 +188,7 @@ run {
 task release(type: Copy, group: 'build') {
     dependsOn allprojects*.tasks.build
     from(zipTree(project.tasks.bundlePlugin.outputs.files.getSingleFile()))
-    into "build/plugins/opensearch-machine-learning"
+    into "build/plugins/opensearch-ml"
     includeEmptyDirs = false
 }
 

plugin/src/test/java/org/opensearch/ml/rest/MLCommonsRestTestCase.java

@@ -18,11 +18,14 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
 import org.apache.http.Header;
@@ -51,17 +54,23 @@
 import org.opensearch.common.xcontent.XContentParser;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.commons.rest.SecureRestClientBuilder;
+import org.opensearch.ml.common.dataset.MLInputDataset;
+import org.opensearch.ml.common.dataset.SearchQueryInputDataset;
 import org.opensearch.ml.common.parameter.FunctionName;
+import org.opensearch.ml.common.parameter.MLAlgoParams;
+import org.opensearch.ml.common.parameter.MLInput;
 import org.opensearch.ml.stats.ActionName;
 import org.opensearch.ml.stats.StatNames;
 import org.opensearch.ml.utils.TestData;
 import org.opensearch.ml.utils.TestHelper;
 import org.opensearch.rest.RestStatus;
+import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.test.rest.OpenSearchRestTestCase;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.gson.Gson;
+import com.google.gson.JsonArray;
 
 public abstract class MLCommonsRestTestCase extends OpenSearchRestTestCase {
     protected Gson gson = new Gson();
@@ -241,7 +250,7 @@ protected Response ingestIrisData(String indexName) throws IOException {
                 "POST",
                 "_bulk?refresh=true",
                 null,
-                TestHelper.toHttpEntity(TestData.IRIS_DATA),
+                TestHelper.toHttpEntity(TestData.IRIS_DATA.replaceAll("iris_data", indexName)),
                 ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, ""))
             );
         assertEquals(RestStatus.OK, TestHelper.restStatus(statsResponse));
@@ -253,7 +262,7 @@ protected void validateStats(
         ActionName actionName,
         int expectedTotalFailureCount,
         int expectedTotalAlgoFailureCount,
-        int expectedTotalRequestCount,
+        int expectedMinumnTotalRequestCount,
         int expectedTotalAlgoRequestCount
     ) throws IOException {
         Response statsResponse = TestHelper.makeRequest(client(), "GET", "_plugins/_ml/stats", null, "", null);
@@ -284,8 +293,7 @@ protected void validateStats(
         }
         assertEquals(expectedTotalFailureCount, totalFailureCount);
         assertEquals(expectedTotalAlgoFailureCount, totalAlgoFailureCount);
-        // ToDo: this line makes this test flaky as other tests makes the request count not predictable
-        // assertEquals(expectedTotalRequestCount, totalRequestCount);
+        assertTrue(totalRequestCount >= expectedMinumnTotalRequestCount);
         assertEquals(expectedTotalAlgoRequestCount, totalAlgoRequestCount);
     }
 
@@ -296,4 +304,164 @@ protected Response ingestModelData() throws IOException {
         assertNotNull(trainModelResponse);
         return trainModelResponse;
     }
+
+    public Response createIndexRole(String role, String index) throws IOException {
+        return TestHelper
+            .makeRequest(
+                client(),
+                "PUT",
+                "/_opendistro/_security/api/roles/" + role,
+                null,
+                TestHelper
+                    .toHttpEntity(
+                        "{\n"
+                            + "\"cluster_permissions\": [\n"
+                            + "],\n"
+                            + "\"index_permissions\": [\n"
+                            + "{\n"
+                            + "\"index_patterns\": [\n"
+                            + "\""
+                            + index
+                            + "\"\n"
+                            + "],\n"
+                            + "\"dls\": \"\",\n"
+                            + "\"fls\": [],\n"
+                            + "\"masked_fields\": [],\n"
+                            + "\"allowed_actions\": [\n"
+                            + "\"crud\",\n"
+                            + "\"indices:admin/create\"\n"
+                            + "]\n"
+                            + "}\n"
+                            + "],\n"
+                            + "\"tenant_permissions\": []\n"
+                            + "}"
+                    ),
+                ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+            );
+    }
+
+    public Response createSearchRole(String role, String index) throws IOException {
+        return TestHelper
+            .makeRequest(
+                client(),
+                "PUT",
+                "/_opendistro/_security/api/roles/" + role,
+                null,
+                TestHelper
+                    .toHttpEntity(
+                        "{\n"
+                            + "\"cluster_permissions\": [\n"
+                            + "],\n"
+                            + "\"index_permissions\": [\n"
+                            + "{\n"
+                            + "\"index_patterns\": [\n"
+                            + "\""
+                            + index
+                            + "\"\n"
+                            + "],\n"
+                            + "\"dls\": \"\",\n"
+                            + "\"fls\": [],\n"
+                            + "\"masked_fields\": [],\n"
+                            + "\"allowed_actions\": [\n"
+                            + "\"indices:data/read/search\"\n"
+                            + "]\n"
+                            + "}\n"
+                            + "],\n"
+                            + "\"tenant_permissions\": []\n"
+                            + "}"
+                    ),
+                ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+            );
+    }
+
+    public Response createUser(String name, String password, ArrayList<String> backendRoles) throws IOException {
+        JsonArray backendRolesString = new JsonArray();
+        for (int i = 0; i < backendRoles.size(); i++) {
+            backendRolesString.add(backendRoles.get(i));
+        }
+        return TestHelper
+            .makeRequest(
+                client(),
+                "PUT",
+                "/_opendistro/_security/api/internalusers/" + name,
+                null,
+                TestHelper
+                    .toHttpEntity(
+                        " {\n"
+                            + "\"password\": \""
+                            + password
+                            + "\",\n"
+                            + "\"backend_roles\": "
+                            + backendRolesString
+                            + ",\n"
+                            + "\"attributes\": {\n"
+                            + "}} "
+                    ),
+                ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+            );
+    }
+
+    public Response deleteUser(String user) throws IOException {
+        return TestHelper
+            .makeRequest(
+                client(),
+                "DELETE",
+                "/_opendistro/_security/api/internalusers/" + user,
+                null,
+                "",
+                ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+            );
+    }
+
+    public Response createRoleMapping(String role, ArrayList<String> users) throws IOException {
+        JsonArray usersString = new JsonArray();
+        for (int i = 0; i < users.size(); i++) {
+            usersString.add(users.get(i));
+        }
+        return TestHelper
+            .makeRequest(
+                client(),
+                "PUT",
+                "/_opendistro/_security/api/rolesmapping/" + role,
+                null,
+                TestHelper
+                    .toHttpEntity(
+                        "{\n" + "  \"backend_roles\" : [  ],\n" + "  \"hosts\" : [  ],\n" + "  \"users\" : " + usersString + "\n" + "}"
+                    ),
+                ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+            );
+    }
+
+    public void trainAndPredict(
+        RestClient client,
+        FunctionName functionName,
+        String indexName,
+        MLAlgoParams params,
+        SearchSourceBuilder searchSourceBuilder,
+        Consumer<Map<String, Object>> function
+    ) throws IOException {
+        MLInputDataset inputData = SearchQueryInputDataset
+            .builder()
+            .indices(ImmutableList.of(indexName))
+            .searchSourceBuilder(searchSourceBuilder)
+            .build();
+        MLInput kmeansInput = MLInput.builder().algorithm(functionName).parameters(params).inputDataset(inputData).build();
+        Response response = TestHelper
+            .makeRequest(
+                client,
+                "POST",
+                "/_plugins/_ml/_train_predict/" + functionName.name().toLowerCase(Locale.ROOT),
+                ImmutableMap.of(),
+                TestHelper.toHttpEntity(kmeansInput),
+                null
+            );
+        HttpEntity entity = response.getEntity();
+        assertNotNull(response);
+        String entityString = TestHelper.httpEntityToString(entity);
+        Map map = gson.fromJson(entityString, Map.class);
+        Map<String, Object> predictionResult = (Map<String, Object>) map.get("prediction_result");
+        if (function != null) {
+            function.accept(predictionResult);
+        }
+    }
 }

plugin/src/test/java/org/opensearch/ml/rest/RestMLTrainAndPredictIT.java

@@ -12,6 +12,8 @@
 import java.util.function.Consumer;
 
 import org.apache.http.HttpEntity;
+import org.junit.After;
+import org.junit.Before;
 import org.opensearch.client.Response;
 import org.opensearch.index.query.MatchAllQueryBuilder;
 import org.opensearch.ml.common.dataset.MLInputDataset;
@@ -27,11 +29,20 @@
 import com.google.common.collect.ImmutableMap;
 
 public class RestMLTrainAndPredictIT extends MLCommonsRestTestCase {
-    private String irisIndex = "iris_data";
+    private String irisIndex = "iris_data_train_predict_it";
+
+    @Before
+    public void setup() throws IOException {
+        ingestIrisData(irisIndex);
+    }
+
+    @After
+    public void deleteIndices() throws IOException {
+        deleteIndexWithAdminClient(irisIndex);
+    }
 
     public void testTrainAndPredictKmeans() throws IOException {
         validateStats(FunctionName.KMEANS, ActionName.TRAIN_PREDICT, 0, 0, 0, 0);
-        ingestIrisData(irisIndex);
         trainAndPredictKmeansWithCustomParam();
         validateStats(FunctionName.KMEANS, ActionName.TRAIN_PREDICT, 0, 0, 1, 1);