Addressing PR comments, adding full access tests, fixing create workf…

…low bug

Signed-off-by: Joshua Palis <[email protected]>

.github/workflows/test_security.yml

@@ -18,7 +18,7 @@ jobs:
       matrix:
         java: [11, 17, 21]
 
-    name: Run Secuirty Integration Tests on Linux
+    name: Run Security Integration Tests on Linux
     runs-on: ubuntu-latest
     needs: Get-CI-Image-Tag
     container:
@@ -36,7 +36,7 @@ jobs:
         with:
           distribution: 'temurin'
           java-version: ${{ matrix.java }}
-      - name: Run build
+      - name: Run tests
         # switching the user, as OpenSearch cluster can only be started as root/Administrator on linux-deb/linux-rpm/windows-zip.
         run: |
           chown -R 1000:1000 `pwd`

src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java

@@ -17,6 +17,7 @@
 import org.opensearch.client.Client;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.rest.RestStatus;
@@ -265,13 +266,17 @@ void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionList
             SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder().query(query).size(0).timeout(requestTimeOut);
 
             SearchRequest searchRequest = new SearchRequest(CommonValue.GLOBAL_CONTEXT_INDEX).source(searchSourceBuilder);
-
-            client.search(searchRequest, ActionListener.wrap(searchResponse -> {
-                internalListener.onResponse(searchResponse.getHits().getTotalHits().value < maxWorkflow);
-            }, exception -> {
-                logger.error("Unable to fetch the workflows", exception);
-                internalListener.onFailure(new FlowFrameworkException("Unable to fetch the workflows", RestStatus.BAD_REQUEST));
-            }));
+            try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+                client.search(searchRequest, ActionListener.wrap(searchResponse -> {
+                    internalListener.onResponse(searchResponse.getHits().getTotalHits().value < maxWorkflow);
+                }, exception -> {
+                    logger.error("Unable to fetch the workflows", exception);
+                    internalListener.onFailure(new FlowFrameworkException("Unable to fetch the workflows", RestStatus.BAD_REQUEST));
+                }));
+            } catch (Exception e) {
+                logger.error("Unable to fetch the workflows", e);
+                internalListener.onFailure(new FlowFrameworkException(e.getMessage(), ExceptionsHelper.status(e)));
+            }
         }
     }
 

src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java

@@ -54,6 +54,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
@@ -63,7 +64,8 @@
  */
 public abstract class FlowFrameworkRestTestCase extends OpenSearchRestTestCase {
 
-    private static String FULL_ACCESS_ROLE = "flow_framework_full_access";
+    private static String FLOW_FRAMEWORK_FULL_ACCESS_ROLE = "flow_framework_full_access";
+    private static String ML_COMMONS_FULL_ACCESS_ROLE = "ml_full_access";
     private static String READ_ACCESS_ROLE = "flow_framework_read_access";
     public static String FULL_ACCESS_USER = "fullAccessUser";
     public static String READ_ACCESS_USER = "readAccessUser";
@@ -129,8 +131,12 @@ protected void setUpSettings() throws Exception {
             String fullAccessUserPassword = generatePassword(FULL_ACCESS_USER);
             String readAccessUserPassword = generatePassword(READ_ACCESS_USER);
 
-            // Configure full access user and client
-            Response response = createUser(FULL_ACCESS_USER, fullAccessUserPassword, FULL_ACCESS_ROLE);
+            // Configure full access user and client, needs ML Full Access role as well
+            Response response = createUser(
+                FULL_ACCESS_USER,
+                fullAccessUserPassword,
+                List.of(FLOW_FRAMEWORK_FULL_ACCESS_ROLE, ML_COMMONS_FULL_ACCESS_ROLE)
+            );
             fullAccessClient = new SecureRestClientBuilder(
                 getClusterHosts().toArray(new HttpHost[0]),
                 isHttps(),
@@ -139,7 +145,7 @@ protected void setUpSettings() throws Exception {
             ).setSocketTimeout(60000).build();
 
             // Configure read access user and client
-            response = createUser(READ_ACCESS_USER, readAccessUserPassword, READ_ACCESS_ROLE);
+            response = createUser(READ_ACCESS_USER, readAccessUserPassword, List.of(READ_ACCESS_ROLE));
             readAccessClient = new SecureRestClientBuilder(
                 getClusterHosts().toArray(new HttpHost[0]),
                 isHttps(),
@@ -264,7 +270,7 @@ protected boolean preserveClusterSettings() {
     }
 
     /**
-     * Create an unguessable password. Simple password are weak due to https://tinyurl.com/383em9zk
+     * Create an unique password. Simple password are weak due to https://tinyurl.com/383em9zk
      * @return a random password.
      */
     public static String generatePassword(String username) {
@@ -565,12 +571,13 @@ protected List<ResourceCreated> getResourcesCreated(RestClient client, String wo
         }
     }
 
-    protected Response createUser(String name, String password, String backendRole) throws IOException {
+    protected Response createUser(String name, String password, List<String> backendRoles) throws IOException {
+        String backendRolesString = backendRoles.stream().map(item -> "\"" + item + "\"").collect(Collectors.joining(","));
         String json = "{\"password\": \""
             + password
-            + "\",\"opendistro_security_roles\": [\""
-            + backendRole
-            + "\"],\"backend_roles\": [],\"attributes\": {}}";
+            + "\",\"opendistro_security_roles\": ["
+            + backendRolesString
+            + "],\"backend_roles\": [],\"attributes\": {}}";
         return TestHelpers.makeRequest(
             client(),
             "PUT",

src/test/java/org/opensearch/flowframework/rest/FlowFrameworkSecureRestApiIT.java

@@ -8,6 +8,7 @@
  */
 package org.opensearch.flowframework.rest;
 
+import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
 import org.opensearch.common.util.io.IOUtils;
@@ -18,6 +19,9 @@
 import org.junit.After;
 
 import java.io.IOException;
+import java.util.Map;
+
+import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_ID;
 
 public class FlowFrameworkSecureRestApiIT extends FlowFrameworkRestTestCase {
 
@@ -57,31 +61,71 @@ public void testGetWorkflowStepsWithReadAccess() throws Exception {
     public void testGetWorkflowWithReadAccess() throws Exception {
         // No permissions to create, so we assert only that the response status isnt forbidden
         ResponseException exception = expectThrows(ResponseException.class, () -> getWorkflow(readAccessClient(), "test"));
-        assertTrue(exception.getMessage().contains("There are no templates in the global_context"));
         assertEquals(RestStatus.NOT_FOUND, TestHelpers.restStatus(exception.getResponse()));
     }
 
     public void testSearchWorkflowWithReadAccess() throws Exception {
+        // Use full access client to invoke create workflow to ensure the template/state indices are created
+        Template template = TestHelpers.createTemplateFromFile("createconnector-registerremotemodel-deploymodel.json");
+        Response response = createWorkflow(fullAccessClient(), template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
         // No permissions to create, so we assert only that the response status isnt forbidden
         String termIdQuery = "{\"query\":{\"ids\":{\"values\":[\"test\"]}}}";
-        ResponseException exception = expectThrows(ResponseException.class, () -> searchWorkflows(readAccessClient(), termIdQuery));
-        assertTrue(exception.getMessage().contains("no such index [.plugins-flow-framework-templates]"));
-        assertEquals(RestStatus.NOT_FOUND, TestHelpers.restStatus(exception.getResponse()));
+        SearchResponse seachResponse = searchWorkflows(readAccessClient(), termIdQuery);
+        assertEquals(RestStatus.OK, seachResponse.status());
     }
 
     public void testGetWorkflowStateWithReadAccess() throws Exception {
+        // Use the full access client to invoke create workflow to ensure the template/state indices are created
+        Template template = TestHelpers.createTemplateFromFile("createconnector-registerremotemodel-deploymodel.json");
+        Response response = createWorkflow(fullAccessClient(), template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
         // No permissions to create or provision, so we assert only that the response status isnt forbidden
         ResponseException exception = expectThrows(ResponseException.class, () -> getWorkflowStatus(readAccessClient(), "test", false));
         assertTrue(exception.getMessage().contains("Fail to find workflow"));
         assertEquals(RestStatus.NOT_FOUND, TestHelpers.restStatus(exception.getResponse()));
     }
 
     public void testSearchWorkflowStateWithReadAccess() throws Exception {
+        // Use the full access client to invoke create workflow to ensure the template/state indices are created
+        Template template = TestHelpers.createTemplateFromFile("createconnector-registerremotemodel-deploymodel.json");
+        Response response = createWorkflow(fullAccessClient(), template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
         // No permissions to create, so we assert only that the response status isnt forbidden
         String termIdQuery = "{\"query\":{\"ids\":{\"values\":[\"test\"]}}}";
-        ResponseException exception = expectThrows(ResponseException.class, () -> searchWorkflowState(readAccessClient(), termIdQuery));
-        assertTrue(exception.getMessage().contains("no such index [.plugins-flow-framework-state]"));
-        assertEquals(RestStatus.NOT_FOUND, TestHelpers.restStatus(exception.getResponse()));
+        SearchResponse searchResponse = searchWorkflowState(readAccessClient(), termIdQuery);
+        assertEquals(RestStatus.OK, searchResponse.status());
+    }
+
+    public void testCreateProvisionDeprovisionWorkflowWithFullAccess() throws Exception {
+        // Invoke create workflow API
+        Template template = TestHelpers.createTemplateFromFile("createconnector-registerremotemodel-deploymodel.json");
+        Response response = createWorkflow(fullAccessClient(), template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+
+        // Retrieve workflow ID
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+
+        // Invoke provision API
+        response = provisionWorkflow(fullAccessClient(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
+
+        // Invoke status API
+        response = getWorkflowStatus(fullAccessClient(), workflowId, false);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
+
+        // Invoke deprovision API
+        response = deprovisionWorkflow(fullAccessClient(), workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
+    }
+
+    public void testGetWorkflowStepWithFullAccess() throws Exception {
+        Response response = getWorkflowStep(fullAccessClient());
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
     }
 
 }