Update Replace.java

Author: sammynorm

kafka-ui-api/src/main/java/com/provectus/kafka/ui/service/masking/policies/Replace.java

@@ -5,17 +5,17 @@
 import com.fasterxml.jackson.databind.node.ContainerNode;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.fasterxml.jackson.databind.node.TextNode;
-import com.google.common.base.Preconditions;
+import java.util.Set;
 
 class Replace extends MaskingPolicy {
 
-  static final String DEFAULT_REPLACEMENT = "***DATA_MASKED***";
-
   private final String replacement;
 
-  Replace(FieldsSelector fieldsSelector, String replacementString) {
+  static final String DEFAULT_REPLACEMENT = "ANONYMIZED";
+
+  Replace(FieldsSelector fieldsSelector, String replacement) {
     super(fieldsSelector);
-    this.replacement = Preconditions.checkNotNull(replacementString);
+    this.replacement = replacement;
   }
 
   @Override
@@ -24,42 +24,29 @@ public String applyToString(String str) {
   }
 
   @Override
-  public ContainerNode<?> applyToJsonContainer(ContainerNode<?> node) {
-    return (ContainerNode<?>) replaceWithFieldsCheck(node);
+  public ContainerNode<?> applyToJsonContainer(ContainerNode<?> node, Set<String> maskedFields) {
+    return (ContainerNode<?>) replaceFields(node, maskedFields);
   }
 
-  private JsonNode replaceWithFieldsCheck(JsonNode node) {
+  private JsonNode replaceFields(JsonNode node, Set<String> maskedFields) {
     if (node.isObject()) {
       ObjectNode obj = ((ObjectNode) node).objectNode();
       node.fields().forEachRemaining(f -> {
         String fieldName = f.getKey();
         JsonNode fieldVal = f.getValue();
         if (fieldShouldBeMasked(fieldName)) {
-          obj.set(fieldName, replaceRecursive(fieldVal));
+          maskedFields.add(fieldName); // ✅ track replaced field
+          obj.set(fieldName, new TextNode(replacement));
         } else {
-          obj.set(fieldName, replaceWithFieldsCheck(fieldVal));
+          obj.set(fieldName, replaceFields(fieldVal, maskedFields));
         }
       });
       return obj;
     } else if (node.isArray()) {
       ArrayNode arr = ((ArrayNode) node).arrayNode(node.size());
-      node.elements().forEachRemaining(e -> arr.add(replaceWithFieldsCheck(e)));
+      node.elements().forEachRemaining(e -> arr.add(replaceFields(e, maskedFields)));
       return arr;
     }
-    // if it is not an object or array - we have nothing to replace here
     return node;
   }
-
-  private JsonNode replaceRecursive(JsonNode node) {
-    if (node.isObject()) {
-      ObjectNode obj = ((ObjectNode) node).objectNode();
-      node.fields().forEachRemaining(f -> obj.set(f.getKey(), replaceRecursive(f.getValue())));
-      return obj;
-    } else if (node.isArray()) {
-      ArrayNode arr = ((ArrayNode) node).arrayNode(node.size());
-      node.elements().forEachRemaining(e -> arr.add(replaceRecursive(e)));
-      return arr;
-    }
-    return new TextNode(replacement);
-  }
 }