diff --git a/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/GetWorkflowTransportAction.java index 11c135210..2eb8b2e39 100644 --- 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 { diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index df7e66e07..551ab7c9b 100644 --- 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 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 processedWorkflows = new HashMap<>(); - for (Map.Entry entry : template.workflows().entrySet()) { + + for (Map.Entry entry : updatedTemplate.workflows().entrySet()) { List 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 index cae595430..c6ec15a92 100644 --- 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()); + } }