[WIP]process-state: filter out additional contents from state downloads

agent-operator/pom.xml

@@ -39,13 +39,16 @@
         </dependency>
         <dependency>
             <groupId>org.bouncycastle</groupId>
-            <artifactId>bcprov-ext-jdk18on</artifactId>
+            <artifactId>bcprov-jdk18on</artifactId>
         </dependency>
         <dependency>
             <groupId>org.bouncycastle</groupId>
             <artifactId>bcpkix-jdk18on</artifactId>
         </dependency>
-
+        <dependency>
+            <groupId>org.bouncycastle</groupId>
+            <artifactId>bcutil-jdk18on</artifactId>
+        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-core</artifactId>

server/impl/src/main/java/com/walmartlabs/concord/server/process/ProcessResource.java

@@ -747,15 +747,17 @@ public void appendLog(@PathParam("id") UUID instanceId, InputStream data) {
             content = @Content(mediaType = "application/zip",
                     schema = @Schema(type = "string", format = "binary"))
     )
-    public Response downloadState(@PathParam("id") UUID instanceId) {
+    public Response downloadState(@PathParam("id") UUID instanceId, @Context HttpServletRequest request) {
         ProcessEntry entry = assertProcess(PartialProcessKey.from(instanceId));
         ProcessKey processKey = new ProcessKey(entry.instanceId(), entry.createdAt());
 
         assertProcessAccess(entry, "state");
 
+        boolean applyFilter = assertApplyFilter(request);
+
         StreamingOutput out = output -> {
             try (ZipArchiveOutputStream dst = new ZipArchiveOutputStream(output)) {
-                stateManager.export(processKey, new ProcessStateManager.FilteringConsumer(zipTo(dst), s -> {
+                stateManager.export(processKey, new ProcessStateManager.FilteringConsumer(zipTo(dst, applyFilter), s -> {
                     if (!isSessionResource(s)) {
                         return true;
                     }
@@ -769,6 +771,21 @@ public Response downloadState(@PathParam("id") UUID instanceId) {
                 .build();
     }
 
+    private boolean assertApplyFilter(HttpServletRequest request) {
+        // do not filter request from admins
+        if (Roles.isAdmin()) {
+            return false;
+        }
+
+        // do not filter request from agent
+        String userAgent = request.getHeader("User-Agent");
+        if (userAgent.startsWith("Concord-Agent")) {
+            return false;
+        }
+
+        return true;
+    }
+
     /**
      * Downloads a single file from the current state snapshot of a process.
      */

server/impl/src/main/java/com/walmartlabs/concord/server/process/StateManagerUtils.java

@@ -0,0 +1,67 @@
+package com.walmartlabs.concord.server.process;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static com.walmartlabs.concord.sdk.Constants.Files.CONFIGURATION_FILE_NAME;
+
+public final class StateManagerUtils {
+
+    private static final Map<String, List<String>> STATE_FILTER = Map.of(CONFIGURATION_FILE_NAME, List.of("arguments"));
+    private static final List<String> ALLOWED_EXTENSIONS = List.of("json", "yaml", "yml");
+
+    public static InputStream stateFilter(String file, InputStream in) {
+        try {
+            String extension = getFileExtension(file);
+            if (!ALLOWED_EXTENSIONS.contains(extension) || STATE_FILTER.get(file) == null) {
+                // only filter for allowed extension files
+                return in;
+            }
+
+            byte[] inputBytes = in.readAllBytes();
+            if (inputBytes.length == 0) {
+                return new ByteArrayInputStream(inputBytes);
+            }
+
+            Map<String, Object> map = switch (extension) {
+                case "json" -> new ObjectMapper().readValue(inputBytes, Map.class);
+                case "yaml", "yml" -> new ObjectMapper(new YAMLFactory()).readValue(inputBytes, Map.class);
+                default -> null;
+            };
+
+            if (map == null) {
+                return new ByteArrayInputStream(inputBytes);
+            }
+
+            Map<String, Object> filteredMap = STATE_FILTER.get(file).stream()
+                    .filter(map::containsKey)
+                    .collect(HashMap::new, (m, key) -> m.put(key, map.get(key)), HashMap::putAll);
+
+            byte[] data = switch (extension) {
+                case "json" -> new ObjectMapper().writeValueAsBytes(filteredMap);
+                case "yaml", "yml" -> new ObjectMapper(new YAMLFactory()).writeValueAsBytes(filteredMap);
+                default -> throw new IllegalArgumentException("Unsupported file extension: " + extension);
+            };
+
+            return new ByteArrayInputStream(data);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static String getFileExtension(String fileName) {
+        if (fileName.lastIndexOf(".") != -1 && fileName.lastIndexOf(".") != 0)
+            return fileName.substring(fileName.lastIndexOf(".") + 1);
+        else return "";
+    }
+
+    private StateManagerUtils() {
+    }
+}

server/impl/src/main/java/com/walmartlabs/concord/server/process/state/ProcessStateManager.java

@@ -34,6 +34,7 @@
 import com.walmartlabs.concord.server.cfg.SecretStoreConfiguration;
 import com.walmartlabs.concord.server.policy.PolicyException;
 import com.walmartlabs.concord.server.policy.PolicyManager;
+import com.walmartlabs.concord.server.process.StateManagerUtils;
 import com.walmartlabs.concord.server.process.logs.ProcessLogManager;
 import com.walmartlabs.concord.server.sdk.PartialProcessKey;
 import com.walmartlabs.concord.server.sdk.ProcessKey;
@@ -510,8 +511,8 @@ public static ItemConsumer copyTo(Path dst, String[] ignored, OpenOption... opti
      *
      * @param dst archive stream.
      */
-    public static ItemConsumer zipTo(ZipArchiveOutputStream dst) {
-        return new ZipConsumer(dst);
+    public static ItemConsumer zipTo(ZipArchiveOutputStream dst, boolean filterContents) {
+        return new ZipConsumer(dst, filterContents);
     }
 
     public static ItemConsumer exclude(ItemConsumer delegate, String... patterns) {
@@ -792,19 +793,22 @@ public void accept(String name, int unixMode, InputStream src) {
     public static final class ZipConsumer implements ItemConsumer {
 
         private final ZipArchiveOutputStream dst;
+        private final boolean filterContents;
 
-        private ZipConsumer(ZipArchiveOutputStream dst) {
+        private ZipConsumer(ZipArchiveOutputStream dst, boolean filterContents) {
             this.dst = dst;
+            this.filterContents = filterContents;
         }
 
         @Override
         public void accept(String name, int unixMode, InputStream src) {
             ZipArchiveEntry entry = new ZipArchiveEntry(name);
             entry.setUnixMode(unixMode);
-
+            // filter before zip to download
+            InputStream processed = filterContents ? StateManagerUtils.stateFilter(name, src) : src;
             try {
                 dst.putArchiveEntry(entry);
-                IOUtils.copy(src, dst);
+                IOUtils.copy(processed, dst);
                 dst.closeArchiveEntry();
             } catch (IOException e) {
                 throw new RuntimeException(e);

server/impl/src/test/java/com/walmartlabs/concord/server/process/StateManagerUtilsTest.java

@@ -0,0 +1,90 @@
+package com.walmartlabs.concord.server.process;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+public class StateManagerUtilsTest {
+
+    @Test
+    void testFilterWithJsonFile() throws Exception {
+        String jsonInput = "{\"arguments\": \"value\", \"otherKey\": \"otherValue\"}";
+        InputStream inputStream = new ByteArrayInputStream(jsonInput.getBytes(StandardCharsets.UTF_8));
+
+        InputStream result = StateManagerUtils.stateFilter("_main.json", inputStream);
+        String text = new String(result.readAllBytes(), StandardCharsets.UTF_8);
+
+        ObjectMapper mapper = new ObjectMapper();
+        Map<String, Object> resultMap = mapper.readValue(text, Map.class);
+
+        assertEquals(1, resultMap.size());
+        assertEquals("value", resultMap.get("arguments"));
+    }
+
+    @Test
+    void testFilterWithYamlFile() throws Exception {
+        String yamlInput = "arguments: value\notherKey: otherValue";
+        InputStream inputStream = new ByteArrayInputStream(yamlInput.getBytes(StandardCharsets.UTF_8));
+
+        InputStream result = StateManagerUtils.stateFilter("test.yaml", inputStream);
+
+        ObjectMapper mapper = new ObjectMapper(new YAMLFactory());
+        Map<String, Object> resultMap = mapper.readValue(result, Map.class);
+
+        // return the same if to stateFilter items not present
+        assertEquals(2, resultMap.size());
+        assertEquals("otherValue", resultMap.get("otherKey"));
+    }
+
+    @Test
+    void testFilterWithUnsupportedExtension() throws Exception {
+        String input = "key: value";
+        InputStream inputStream = new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8));
+
+        InputStream result = StateManagerUtils.stateFilter("file.txt", inputStream);
+        String text = new String(result.readAllBytes(), StandardCharsets.UTF_8);
+
+        assertEquals(input, text);
+    }
+
+    // Test case for empty input
+    @Test
+    void testFilterWithEmptyInput() throws Exception {
+        String input = "";
+        InputStream inputStream = new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8));
+
+        InputStream result = StateManagerUtils.stateFilter("_main.json", inputStream);
+        String text = new String(result.readAllBytes(), StandardCharsets.UTF_8);
+
+        assertEquals(input, text);
+    }
+
+    // Test case for null filtering rules
+    @Test
+    void testFilterWithNullRules() throws Exception {
+        String jsonInput = "{\"key\": \"value\"}";
+        InputStream inputStream = new ByteArrayInputStream(jsonInput.getBytes(StandardCharsets.UTF_8));
+
+        InputStream result = StateManagerUtils.stateFilter("unknown.json", inputStream);
+        String text = new String(result.readAllBytes(), StandardCharsets.UTF_8);
+
+        assertEquals(jsonInput, text);
+    }
+
+    // Test case for null input
+    @Test
+    void testFilterWithNullInput() throws Exception {
+        String jsonInput = "{\"key\": \"value\"}";
+        InputStream inputStream = null;
+        InputStream result = StateManagerUtils.stateFilter("unknown.json", inputStream);
+
+        assertNull(result);
+    }
+}

targetplatform/pom.xml

@@ -39,7 +39,6 @@
         <aopalliance.version>1.0</aopalliance.version>
         <aws.version>1.11.475</aws.version>
         <bouncycastle.version>1.80</bouncycastle.version>
-        <bouncycastle.ext.version>1.78.1</bouncycastle.ext.version>
         <bpm.version>1.0.3</bpm.version>
         <bytebuddy.version>1.14.9</bytebuddy.version>
         <commons.beanutils.version>1.9.4</commons.beanutils.version>
@@ -1103,17 +1102,17 @@
             </dependency>
             <dependency>
                 <groupId>org.bouncycastle</groupId>
-                <artifactId>bcprov-ext-jdk18on</artifactId>
-                <version>${bouncycastle.ext.version}</version>
+                <artifactId>bcprov-jdk18on</artifactId>
+                <version>${bouncycastle.version}</version>
             </dependency>
             <dependency>
                 <groupId>org.bouncycastle</groupId>
-                <artifactId>bcprov-jdk18on</artifactId>
+                <artifactId>bcpkix-jdk18on</artifactId>
                 <version>${bouncycastle.version}</version>
             </dependency>
             <dependency>
                 <groupId>org.bouncycastle</groupId>
-                <artifactId>bcpkix-jdk18on</artifactId>
+                <artifactId>bcutil-jdk18on</artifactId>
                 <version>${bouncycastle.version}</version>
             </dependency>
             <dependency>