runtime-v2: use SensitiveDataHolder for task parameter masking

Author: ibodrov

runtime/v2/runner/src/main/java/com/walmartlabs/concord/runtime/v2/runner/remote/TaskCallEventRecordingListener.java

@@ -28,6 +28,7 @@
 import com.walmartlabs.concord.runtime.v2.model.Location;
 import com.walmartlabs.concord.runtime.v2.model.Step;
 import com.walmartlabs.concord.runtime.v2.runner.EventReportingService;
+import com.walmartlabs.concord.runtime.v2.runner.SensitiveDataHolder;
 import com.walmartlabs.concord.runtime.v2.runner.tasks.TaskCallEvent;
 import com.walmartlabs.concord.runtime.v2.runner.tasks.TaskCallListener;
 import com.walmartlabs.concord.runtime.v2.sdk.*;
@@ -64,7 +65,9 @@ public void onEvent(TaskCallEvent event) {
 
         List<Object> inVars = event.input();
         if (inVars != null && eventConfiguration.recordTaskInVars()) {
-            Map<String, Object> vars = maskVars(convertInput(hideSensitiveData(inVars, event.inputAnnotations())), eventConfiguration.inVarsBlacklist());
+            Map<String, Object> input = convertInput(processSensitiveDataAnnotations(inVars, event.inputAnnotations()));
+            input = processSensitiveData(input);
+            Map<String, Object> vars = processVarBlacklist(input, eventConfiguration.inVarsBlacklist());
             if (eventConfiguration.truncateInVars()) {
                 vars = ObjectTruncater.truncateMap(vars, eventConfiguration.truncateMaxStringLength(), eventConfiguration.truncateMaxArrayLength(), eventConfiguration.truncateMaxDepth());
             }
@@ -75,7 +78,9 @@ public void onEvent(TaskCallEvent event) {
 
         Object outVars = event.result();
         if (outVars != null && eventConfiguration.recordTaskOutVars()) {
-            Map<String, Object> vars = maskVars(asMapOrNull(outVars), eventConfiguration.outVarsBlacklist());
+            Map<String, Object> output = asMapOrNull(outVars);
+            output = processSensitiveData(output);
+            Map<String, Object> vars = processVarBlacklist(output, eventConfiguration.outVarsBlacklist());
             if (eventConfiguration.truncateOutVars()) {
                 vars = ObjectTruncater.truncateMap(vars, eventConfiguration.truncateMaxStringLength(), eventConfiguration.truncateMaxArrayLength(), eventConfiguration.truncateMaxDepth());
             }
@@ -86,7 +91,9 @@ public void onEvent(TaskCallEvent event) {
 
         Object metaVars = event.meta();
         if (metaVars != null && eventConfiguration.recordTaskMeta()) {
-            Map<String, Object> meta = maskVars(asMapOrNull(metaVars), eventConfiguration.metaBlacklist());
+            Map<String, Object> rawMeta = asMapOrNull(metaVars);
+            Map<String, Object> meta = processSensitiveData(rawMeta);
+            meta = processVarBlacklist(meta, eventConfiguration.metaBlacklist());
             if (eventConfiguration.truncateMeta()) {
                 meta = ObjectTruncater.truncateMap(meta, eventConfiguration.truncateMaxStringLength(), eventConfiguration.truncateMaxArrayLength(), eventConfiguration.truncateMaxDepth());
             }
@@ -154,7 +161,7 @@ private Map<String, Object> asMapOrNull(Object v) {
         return null;
     }
 
-    static Map<String, Object> maskVars(Map<String, Object> vars, Collection<String> blackList) {
+    static Map<String, Object> processVarBlacklist(Map<String, Object> vars, Collection<String> blackList) {
         if (blackList.isEmpty()) {
             return vars;
         }
@@ -170,6 +177,34 @@ static Map<String, Object> maskVars(Map<String, Object> vars, Collection<String>
         return result;
     }
 
+    @SuppressWarnings({"unchecked", "rawtypes"})
+    static <T> T processSensitiveData(T v) {
+        Set<String> sensitiveStrings = SensitiveDataHolder.getInstance().get();
+        if (sensitiveStrings.isEmpty()) {
+            return v;
+        }
+
+        if (v instanceof String s) {
+            for (String sensitiveString : sensitiveStrings) {
+                s = s.replace(sensitiveString, MASK);
+            }
+            return (T) s;
+        } else if (v instanceof List<?> l) {
+            List<Object> result = new ArrayList<>(l.size());
+            for (Object vv : l) {
+                vv = processSensitiveData(vv);
+                result.add(vv);
+            }
+            return (T) result;
+        } else if (v instanceof Map m) {
+            Map<String, Object> result = new HashMap<>(m);
+            result.replaceAll((k, vv) -> processSensitiveData(vv));
+            return (T) result;
+        }
+
+        return v;
+    }
+
     @SuppressWarnings("unchecked")
     private static Map<String, Object> ensureModifiable(Map<String, Object> m, int depth, String[] path) {
         if (depth == 0) {
@@ -217,7 +252,7 @@ private static Map<String, Object> convertInput(List<Object> input) {
         return result;
     }
 
-    private static List<Object> hideSensitiveData(List<Object> input, List<List<Annotation>> annotations) {
+    private static List<Object> processSensitiveDataAnnotations(List<Object> input, List<List<Annotation>> annotations) {
         if (annotations.isEmpty()) {
             return input;
         }

runtime/v2/runner/src/test/java/com/walmartlabs/concord/runtime/v2/runner/remote/TaskCallEventRecordingListenerTest.java

@@ -51,7 +51,7 @@ public void testMaskVars() throws Exception {
                 "}";
 
         List<String> blackList = Arrays.asList("b", "c.c1", "c.c3.c31");
-        Map<String, Object> result = TaskCallEventRecordingListener.maskVars(vars(in), blackList);
+        Map<String, Object> result = TaskCallEventRecordingListener.processVarBlacklist(vars(in), blackList);
 
         String expected = "{" +
                 "   \"a\":1," +
@@ -76,7 +76,7 @@ public void testMaskVarsUnmodifiable() {
                                 Collections.singletonMap("z", 123)));
 
         List<String> blackList = Collections.singletonList("x.y.z");
-        Map<String, Object> result = TaskCallEventRecordingListener.maskVars(vars, blackList);
+        Map<String, Object> result = TaskCallEventRecordingListener.processVarBlacklist(vars, blackList);
         assertEquals("{x={y={z=***}}}", result.toString());
     }