From 359eb809326518fcfdc4e3feb85afc234d59bf31 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Fri, 12 Jan 2024 22:24:45 +0000 Subject: [PATCH] Addressing PR comments, adding full access tests, fixing create workflow bug Signed-off-by: Joshua Palis --- .github/workflows/test_security.yml | 4 +- .../CreateWorkflowTransportAction.java | 19 +++--- .../FlowFrameworkRestTestCase.java | 25 +++++--- .../rest/FlowFrameworkSecureRestApiIT.java | 58 ++++++++++++++++--- 4 files changed, 81 insertions(+), 25 deletions(-) diff --git a/.github/workflows/test_security.yml b/.github/workflows/test_security.yml index 373d857ed..c18b2a11a 100644 --- a/.github/workflows/test_security.yml +++ b/.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` diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 2fd9bd042..2c80969d6 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/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))); + } } } diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index ff64fdc57..1d450cd26 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/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 getResourcesCreated(RestClient client, String wo } } - protected Response createUser(String name, String password, String backendRole) throws IOException { + protected Response createUser(String name, String password, List 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", diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkSecureRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkSecureRestApiIT.java index f0ffb63d9..44fed3b5b 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkSecureRestApiIT.java +++ b/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,19 +61,27 @@ 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")); @@ -77,11 +89,43 @@ public void testGetWorkflowStateWithReadAccess() throws Exception { } 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 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)); } }