Hid the user field for Get API as well

Signed-off-by: Owais Kazi <[email protected]>
opensearch-project · Apr 25, 2024 · 35f04e8 · 35f04e8
1 parent 6228342
commit 35f04e8
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 6 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java
@@ -80,7 +80,7 @@ protected void doExecute(Task task, WorkflowRequest request, ActionListener<GetW
                         listener.onFailure(new FlowFrameworkException(errorMessage, RestStatus.NOT_FOUND));
                     } else {
                         // Remove any credential from response
-                        Template template = encryptorUtils.redactTemplateCredentials(Template.parse(response.getSourceAsString()));
+                        Template template = encryptorUtils.redactTemplateSecuredFields(Template.parse(response.getSourceAsString()));
                         listener.onResponse(new GetWorkflowResponse(template));
                     }
                 }, exception -> {

diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -122,7 +122,7 @@ public Template decryptTemplateCredentials(Template template) {
     /**
      * Applies the given cipher function on template credentials
      * @param template the template to process
-     * @param cipher the encryption/decryption function to apply on credential values
+     * @param cipherFunction the encryption/decryption function to apply on credential values
      * @return template with encrypted credentials
      */
     private Template processTemplateCredentials(Template template, Function<String, String> cipherFunction) {
@@ -204,9 +204,29 @@ String decrypt(final String encryptedCredential) {
      * @param template the template
      * @return the redacted template
      */
-    public Template redactTemplateCredentials(Template template) {
+    public Template redactTemplateSecuredFields(Template template) {
+        Template updatedTemplate = null;
+
+        if (template.getUser() != null) {
+            updatedTemplate = new Template.Builder(template).name(template.name())
+                .description(template.description())
+                .useCase(template.useCase())
+                .templateVersion(template.templateVersion())
+                .user(null)
+                .uiMetadata(template.getUiMetadata())
+                .compatibilityVersion(template.compatibilityVersion())
+                .workflows(template.workflows())
+                .createdTime(template.createdTime())
+                .lastUpdatedTime(template.lastUpdatedTime())
+                .lastProvisionedTime(template.lastProvisionedTime())
+                .build();
+        } else {
+            updatedTemplate = template;
+        }
+
         Map<String, Workflow> processedWorkflows = new HashMap<>();
-        for (Map.Entry<String, Workflow> entry : template.workflows().entrySet()) {
+
+        for (Map.Entry<String, Workflow> entry : updatedTemplate.workflows().entrySet()) {
 
             List<WorkflowNode> processedNodes = new ArrayList<>();
             for (WorkflowNode node : entry.getValue().nodes()) {
@@ -227,7 +247,7 @@ public Template redactTemplateCredentials(Template template) {
             processedWorkflows.put(entry.getKey(), new Workflow(entry.getValue().userParams(), processedNodes, entry.getValue().edges()));
         }
 
-        return new Template.Builder(template).workflows(processedWorkflows).build();
+        return new Template.Builder(updatedTemplate).workflows(processedWorkflows).build();
     }
 
     /**

diff --git a/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java b/src/test/java/org/opensearch/flowframework/util/EncryptorUtilsTests.java
@@ -200,10 +200,21 @@ public void testRedactTemplateCredential() {
         assertNotNull(node.userInputs().get(CREDENTIAL_FIELD));
 
         // Redact template with credential field
-        Template redactedTemplate = encryptorUtils.redactTemplateCredentials(testTemplate);
+        Template redactedTemplate = encryptorUtils.redactTemplateSecuredFields(testTemplate);
 
         // Validate the credential field has been removed
         WorkflowNode redactedNode = redactedTemplate.workflows().get("provision").nodes().get(0);
         assertNull(redactedNode.userInputs().get(CREDENTIAL_FIELD));
     }
+
+    public void testRedactTemplateUserField() {
+        // Confirm user is present in the non-redacted template
+        assertNotNull(testTemplate.getUser());
+
+        // Redact template with user field
+        Template redactedTemplate = encryptorUtils.redactTemplateSecuredFields(testTemplate);
+
+        // Validate the user field has been removed
+        assertNull(redactedTemplate.getUser());
+    }
 }