fix(masking): Various performance, functional, and readability fixes for JSON-based masking

Author: jamfor999

docs/docs/configuration/akhq.md

@@ -121,6 +121,7 @@ akhq:
             - dateOfBirth
             - address.firstLine
             - address.town
+            - metadata.notes
 ```
 
 Given a record on `users` that looks like:
@@ -131,11 +132,18 @@ Given a record on `users` that looks like:
   "status": "ACTIVE",
   "name": "John Smith",
   "dateOfBirth": "01-01-1991",
-  "address": {
-    "firstLine": "123 Example Avenue",
-    "town": "Faketown",
-    "country": "United Kingdom"
-  },
+  "address": [
+    {
+      "firstLine": "123 Example Avenue",
+      "town": "Faketown",
+      "country": "United Kingdom"
+    },
+    {
+      "firstLine": "123 Previous Avenue",
+      "town": "Previoustown",
+      "country": "United Kingdom"
+    }
+  ],
   "metadata": {
     "trusted": true,
     "rating": "10",
@@ -152,19 +160,51 @@ With the above configuration, it will appear as:
   "status": "ACTIVE",
   "name": "xxxx",
   "dateOfBirth": "xxxx",
-  "address": {
-    "firstLine": "xxxx",
-    "town": "xxxx",
-    "country": "United Kingdom"
-  },
+  "address": [
+    {
+      "firstLine": "xxxx",
+      "town": "xxxx",
+      "country": "United Kingdom"
+    },
+    {
+      "firstLine": "xxxx",
+      "town": "xxxx",
+      "country": "United Kingdom"
+    }
+  ],
   "metadata": {
     "trusted": true,
     "rating": "10",
-    "notes": "All in good order"
+    "notes": "xxxx"
+  }
+}
+```
+
+Note how arrays are automatically understood where relevant.
+In other words, `address.firstLine` will apply to both of the following:
+```json
+{
+  "address": {
+    "firstLine": "This field!"
   }
 }
 ```
 
+and
+
+```json
+{
+  "address": [
+    {
+      "firstLine": "This field!"
+    },
+    {
+      "firstLine": "And this one!"
+    }
+  ]
+}
+```
+
 ### Mask by default config
 This means, by default, everything is masked.
 This is useful in production scenarios where data must be carefully selected and made available to
@@ -207,11 +247,18 @@ Given a record on `users` that looks like:
   "status": "ACTIVE",
   "name": "John Smith",
   "dateOfBirth": "01-01-1991",
-  "address": {
-    "firstLine": "123 Example Avenue",
-    "town": "Faketown",
-    "country": "United Kingdom"
-  },
+  "address": [
+    {
+      "firstLine": "123 Example Avenue",
+      "town": "Faketown",
+      "country": "United Kingdom"
+    },
+    {
+      "firstLine": "123 Previous Avenue",
+      "town": "Previoustown",
+      "country": "United Kingdom"
+    }
+  ],
   "metadata": {
     "trusted": true,
     "rating": "10",
@@ -228,11 +275,18 @@ With the above configuration, it will appear as:
   "status": "ACTIVE",
   "name": "xxxx",
   "dateOfBirth": "xxxx",
-  "address": {
-    "firstLine": "xxxx",
-    "town": "xxxx",
-    "country": "United Kingdom"
-  },
+  "address": [
+    {
+      "firstLine": "xxxx",
+      "town": "xxxx",
+      "country": "United Kingdom"
+    },
+    {
+      "firstLine": "xxxx",
+      "town": "xxxx",
+      "country": "United Kingdom"
+    }
+  ],
   "metadata": {
     "trusted": true,
     "rating": "10",
@@ -241,6 +295,31 @@ With the above configuration, it will appear as:
 }
 ```
 
+Note how arrays are automatically understood where relevant.
+In other words, `address.firstLine` will apply to both of the following:
+```json
+{
+  "address": {
+    "firstLine": "This field!"
+  }
+}
+```
+
+and
+
+```json
+{
+  "address": [
+    {
+      "firstLine": "This field!"
+    },
+    {
+      "firstLine": "And this one!"
+    }
+  ]
+}
+```
+
 ### No masking required
 You can set `akhq.security.data-masking.mode` to `none` to disable masking altogether.
 

src/main/java/org/akhq/utils/JsonMaskByDefaultMasker.java

@@ -1,73 +1,94 @@
 package org.akhq.utils;
 
+import com.google.gson.JsonArray;
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
+import com.google.gson.JsonPrimitive;
 import io.micronaut.context.annotation.Requires;
 import jakarta.inject.Singleton;
 import lombok.SneakyThrows;
 import org.akhq.configs.DataMasking;
-import org.akhq.configs.JsonMaskingFilter;
 import org.akhq.models.Record;
 
 import java.util.List;
 import java.util.Map;
 
 @Singleton
 @Requires(property = "akhq.security.data-masking.mode", value = "json_mask_by_default")
-public class JsonMaskByDefaultMasker implements Masker {
+public class JsonMaskByDefaultMasker extends JsonMasker {
 
-    private final List<JsonMaskingFilter> jsonMaskingFilters;
-    private final String jsonMaskReplacement;
+    private static final String NON_JSON_MESSAGE = "This record is unable to be masked as it is not a structured object. This record is unavailable to view due to safety measures from json_mask_by_default to not leak sensitive data.";
+    private static final String ERROR_MESSAGE = "An exception occurred during an attempt to mask this record. This record is unavailable to view due to safety measures from json_mask_by_default to not leak sensitive data.";
 
     public JsonMaskByDefaultMasker(DataMasking dataMasking) {
-        this.jsonMaskingFilters = dataMasking.getJsonFilters();
-        this.jsonMaskReplacement = dataMasking.getJsonMaskReplacement();
+        super(dataMasking);
     }
 
     public Record maskRecord(Record record) {
+        if (!isJson(record)) {
+            record.setValue(NON_JSON_MESSAGE);
+            return record;
+        }
+
         try {
-            if(isJson(record)) {
-                return jsonMaskingFilters
-                    .stream()
-                    .filter(jsonMaskingFilter -> record.getTopic().getName().equalsIgnoreCase(jsonMaskingFilter.getTopic()))
-                    .findFirst()
-                    .map(filter -> applyMasking(record, filter.getKeys()))
-                    .orElseGet(() -> applyMasking(record, List.of()));
-            } else {
-                record.setValue("This record is unable to be masked as it is not a structured object. This record is unavailable to view due to safety measures from json_mask_by_default to not leak sensitive data. Please contact akhq administrator.");
-            }
+            List<String> keysToUnmask = getKeysForTopic(record.getTopic().getName());
+            return applyMasking(record, keysToUnmask);
         } catch (Exception e) {
-            LOG.error("Error masking record at topic {}, partition {}, offset {} due to {}", record.getTopic(), record.getPartition(), record.getOffset(), e.getMessage());
-            record.setValue("An exception occurred during an attempt to mask this record. This record is unavailable to view due to safety measures from json_mask_by_default to not leak sensitive data. Please contact akhq administrator.");
+            LOG.error("Error masking record at topic {}, partition {}, offset {} due to {}",
+                record.getTopic(), record.getPartition(), record.getOffset(), e.getMessage());
+            record.setValue(ERROR_MESSAGE);
+            return record;
         }
-        return record;
     }
 
     @SneakyThrows
-    private Record applyMasking(Record record, List<String> keys) {
-        JsonObject jsonElement = JsonParser.parseString(record.getValue()).getAsJsonObject();
-        maskAllExcept(jsonElement, keys);
-        record.setValue(jsonElement.toString());
+    private Record applyMasking(Record record, List<String> keysToUnmask) {
+        JsonElement root = JsonParser.parseString(record.getValue());
+        maskJson(root, "", keysToUnmask);
+        record.setValue(root.toString());
         return record;
     }
 
-    private void maskAllExcept(JsonObject jsonElement, List<String> keys) {
-        maskAllExcept("",  jsonElement, keys);
+    private void maskJson(JsonElement element, String path, List<String> keysToUnmask) {
+        if (element.isJsonObject()) {
+            maskJsonObject(element.getAsJsonObject(), path, keysToUnmask);
+        } else if (element.isJsonArray()) {
+            maskJsonArray(element.getAsJsonArray(), path, keysToUnmask);
+        }
     }
 
-    private void maskAllExcept(String currentKey, JsonObject node, List<String> keys) {
-        if (node.isJsonObject()) {
-            JsonObject objectNode = node.getAsJsonObject();
-            for(Map.Entry<String, JsonElement> entry : objectNode.entrySet()) {
-                if(entry.getValue().isJsonObject()) {
-                    maskAllExcept(currentKey + entry.getKey() + ".", entry.getValue().getAsJsonObject(), keys);
-                } else {
-                    if(!keys.contains(currentKey + entry.getKey())) {
-                        objectNode.addProperty(entry.getKey(), jsonMaskReplacement);
-                    }
-                }
+    private void maskJsonObject(JsonObject obj, String path, List<String> keysToUnmask) {
+        for (Map.Entry<String, JsonElement> entry : obj.entrySet()) {
+            String newPath = path + entry.getKey();
+            JsonElement value = entry.getValue();
+
+            if (shouldMaskPrimitive(value, newPath, keysToUnmask)) {
+                entry.setValue(new JsonPrimitive(jsonMaskReplacement));
+            } else if (isNestedStructure(value)) {
+                maskJson(value, newPath + ".", keysToUnmask);
             }
         }
     }
+
+    private void maskJsonArray(JsonArray array, String path, List<String> keysToUnmask) {
+        boolean shouldMask = !keysToUnmask.contains(path.substring(0, path.length() - 1));
+
+        for (int i = 0; i < array.size(); i++) {
+            JsonElement arrayElement = array.get(i);
+            if (arrayElement.isJsonPrimitive() && shouldMask) {
+                array.set(i, new JsonPrimitive(jsonMaskReplacement));
+            } else if (isNestedStructure(arrayElement)) {
+                maskJson(arrayElement, path, keysToUnmask);
+            }
+        }
+    }
+
+    private boolean shouldMaskPrimitive(JsonElement value, String path, List<String> keysToUnmask) {
+        return value.isJsonPrimitive() && !keysToUnmask.contains(path);
+    }
+
+    private boolean isNestedStructure(JsonElement value) {
+        return value.isJsonObject() || value.isJsonArray();
+    }
 }

src/main/java/org/akhq/utils/JsonMasker.java

@@ -0,0 +1,34 @@
+package org.akhq.utils;
+
+import org.akhq.configs.DataMasking;
+import org.akhq.configs.JsonMaskingFilter;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public abstract class JsonMasker implements Masker {
+    private final Map<String, List<String>> topicToKeysMap;
+    protected final String jsonMaskReplacement;
+
+    public JsonMasker(DataMasking dataMasking) {
+        this.jsonMaskReplacement = dataMasking.getJsonMaskReplacement();
+        this.topicToKeysMap = buildTopicKeysMap(dataMasking);
+    }
+
+    private Map<String, List<String>> buildTopicKeysMap(DataMasking dataMasking) {
+        return dataMasking.getJsonFilters().stream()
+            .collect(Collectors.toMap(
+                JsonMaskingFilter::getTopic,
+                JsonMaskingFilter::getKeys,
+                (a, b) -> a,
+                HashMap::new
+            ));
+    }
+
+    protected List<String> getKeysForTopic(String topic) {
+        return topicToKeysMap.getOrDefault(topic.toLowerCase(), Collections.emptyList());
+    }
+}