Skip to content

Commit 657da6d

Browse files
jamfor999jamfor999
committed
fix(masking): Various performance and readability fixes
1 parent 6b43d40 commit 657da6d

File tree

8 files changed

+354
-321
lines changed

8 files changed

+354
-321
lines changed

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

Lines changed: 84 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,73 +1,123 @@
11
package org.akhq.utils;
22

3-
import com.google.gson.JsonElement;
4-
import com.google.gson.JsonObject;
5-
import com.google.gson.JsonParser;
3+
import com.google.gson.*;
64
import io.micronaut.context.annotation.Requires;
75
import jakarta.inject.Singleton;
86
import lombok.SneakyThrows;
97
import org.akhq.configs.DataMasking;
108
import org.akhq.configs.JsonMaskingFilter;
119
import org.akhq.models.Record;
1210

11+
import java.util.Collections;
12+
import java.util.HashMap;
1313
import java.util.List;
1414
import java.util.Map;
15+
import java.util.stream.Collectors;
1516

1617
@Singleton
1718
@Requires(property = "akhq.security.data-masking.mode", value = "json_mask_by_default")
1819
public class JsonMaskByDefaultMasker implements Masker {
1920

20-
private final List<JsonMaskingFilter> jsonMaskingFilters;
21+
private final Map<String, List<String>> topicToKeysMap;
2122
private final String jsonMaskReplacement;
23+
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.";
24+
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. Please contact akhq administrator.";
2225

2326
public JsonMaskByDefaultMasker(DataMasking dataMasking) {
24-
this.jsonMaskingFilters = dataMasking.getJsonFilters();
2527
this.jsonMaskReplacement = dataMasking.getJsonMaskReplacement();
28+
this.topicToKeysMap = buildTopicKeysMap(dataMasking);
29+
}
30+
31+
private Map<String, List<String>> buildTopicKeysMap(DataMasking dataMasking) {
32+
return dataMasking.getJsonFilters().stream()
33+
.collect(Collectors.toMap(
34+
JsonMaskingFilter::getTopic,
35+
JsonMaskingFilter::getKeys,
36+
(a, b) -> a,
37+
HashMap::new
38+
));
2639
}
2740

2841
public Record maskRecord(Record record) {
42+
if (!isJson(record)) {
43+
return createNonJsonRecord(record);
44+
}
45+
2946
try {
30-
if(isJson(record)) {
31-
return jsonMaskingFilters
32-
.stream()
33-
.filter(jsonMaskingFilter -> record.getTopic().getName().equalsIgnoreCase(jsonMaskingFilter.getTopic()))
34-
.findFirst()
35-
.map(filter -> applyMasking(record, filter.getKeys()))
36-
.orElseGet(() -> applyMasking(record, List.of()));
37-
} else {
38-
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.");
39-
}
47+
List<String> unmaskedKeys = getUnmaskedKeysForTopic(record.getTopic().getName());
48+
return applyMasking(record, unmaskedKeys);
4049
} catch (Exception e) {
41-
LOG.error("Error masking record at topic {}, partition {}, offset {} due to {}", record.getTopic(), record.getPartition(), record.getOffset(), e.getMessage());
42-
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.");
50+
logMaskingError(record, e);
51+
return createErrorRecord(record);
4352
}
53+
}
54+
55+
private List<String> getUnmaskedKeysForTopic(String topic) {
56+
return topicToKeysMap.getOrDefault(topic.toLowerCase(), Collections.emptyList());
57+
}
58+
59+
private Record createNonJsonRecord(Record record) {
60+
record.setValue(NON_JSON_MESSAGE);
61+
return record;
62+
}
63+
64+
private Record createErrorRecord(Record record) {
65+
record.setValue(ERROR_MESSAGE);
4466
return record;
4567
}
4668

69+
private void logMaskingError(Record record, Exception e) {
70+
LOG.error("Error masking record at topic {}, partition {}, offset {} due to {}",
71+
record.getTopic(), record.getPartition(), record.getOffset(), e.getMessage());
72+
}
73+
4774
@SneakyThrows
48-
private Record applyMasking(Record record, List<String> keys) {
49-
JsonObject jsonElement = JsonParser.parseString(record.getValue()).getAsJsonObject();
50-
maskAllExcept(jsonElement, keys);
51-
record.setValue(jsonElement.toString());
75+
private Record applyMasking(Record record, List<String> unmaskedKeys) {
76+
JsonObject root = JsonParser.parseString(record.getValue()).getAsJsonObject();
77+
maskJson(root, "", unmaskedKeys);
78+
record.setValue(root.toString());
5279
return record;
5380
}
5481

55-
private void maskAllExcept(JsonObject jsonElement, List<String> keys) {
56-
maskAllExcept("", jsonElement, keys);
82+
private void maskJson(JsonElement element, String path, List<String> unmaskedKeys) {
83+
if (element.isJsonObject()) {
84+
maskJsonObject(element.getAsJsonObject(), path, unmaskedKeys);
85+
} else if (element.isJsonArray()) {
86+
maskJsonArray(element.getAsJsonArray(), path, unmaskedKeys);
87+
}
88+
}
89+
90+
private void maskJsonObject(JsonObject obj, String path, List<String> unmaskedKeys) {
91+
for (Map.Entry<String, JsonElement> entry : obj.entrySet()) {
92+
String newPath = path + entry.getKey();
93+
JsonElement value = entry.getValue();
94+
95+
if (shouldMaskPrimitive(value, newPath, unmaskedKeys)) {
96+
entry.setValue(new JsonPrimitive(jsonMaskReplacement));
97+
} else if (isNestedStructure(value)) {
98+
maskJson(value, newPath + ".", unmaskedKeys);
99+
}
100+
}
57101
}
58102

59-
private void maskAllExcept(String currentKey, JsonObject node, List<String> keys) {
60-
if (node.isJsonObject()) {
61-
JsonObject objectNode = node.getAsJsonObject();
62-
for(Map.Entry<String, JsonElement> entry : objectNode.entrySet()) {
63-
if(entry.getValue().isJsonObject()) {
64-
maskAllExcept(currentKey + entry.getKey() + ".", entry.getValue().getAsJsonObject(), keys);
65-
} else {
66-
if(!keys.contains(currentKey + entry.getKey())) {
67-
objectNode.addProperty(entry.getKey(), jsonMaskReplacement);
68-
}
69-
}
103+
private void maskJsonArray(JsonArray array, String path, List<String> unmaskedKeys) {
104+
boolean shouldMask = !unmaskedKeys.contains(path.substring(0, path.length() - 1));
105+
106+
for (int i = 0; i < array.size(); i++) {
107+
JsonElement arrayElement = array.get(i);
108+
if (arrayElement.isJsonPrimitive() && shouldMask) {
109+
array.set(i, new JsonPrimitive(jsonMaskReplacement));
110+
} else if (isNestedStructure(arrayElement)) {
111+
maskJson(arrayElement, path, unmaskedKeys);
70112
}
71113
}
72114
}
115+
116+
private boolean shouldMaskPrimitive(JsonElement value, String path, List<String> unmaskedKeys) {
117+
return value.isJsonPrimitive() && !unmaskedKeys.contains(path);
118+
}
119+
120+
private boolean isNestedStructure(JsonElement value) {
121+
return value.isJsonObject() || value.isJsonArray();
122+
}
73123
}

‎src/main/java/org/akhq/utils/JsonShowByDefaultMasker.java‎

Lines changed: 86 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,122 @@
11
package org.akhq.utils;
22

3-
import com.google.gson.JsonElement;
4-
import com.google.gson.JsonObject;
5-
import com.google.gson.JsonParser;
3+
import com.google.gson.*;
64
import io.micronaut.context.annotation.Requires;
75
import jakarta.inject.Singleton;
86
import lombok.SneakyThrows;
97
import org.akhq.configs.DataMasking;
108
import org.akhq.configs.JsonMaskingFilter;
119
import org.akhq.models.Record;
1210

11+
import java.util.HashMap;
1312
import java.util.List;
13+
import java.util.Map;
14+
import java.util.stream.Collectors;
1415

1516
@Singleton
1617
@Requires(property = "akhq.security.data-masking.mode", value = "json_show_by_default")
1718
public class JsonShowByDefaultMasker implements Masker {
1819

19-
private final List<JsonMaskingFilter> jsonMaskingFilters;
20+
private final Map<String, List<String>> topicToKeysMap;
2021
private final String jsonMaskReplacement;
22+
private static final String ERROR_MESSAGE = "Error masking record";
2123

2224
public JsonShowByDefaultMasker(DataMasking dataMasking) {
23-
this.jsonMaskingFilters = dataMasking.getJsonFilters();
2425
this.jsonMaskReplacement = dataMasking.getJsonMaskReplacement();
26+
this.topicToKeysMap = buildTopicKeysMap(dataMasking);
27+
}
28+
29+
private Map<String, List<String>> buildTopicKeysMap(DataMasking dataMasking) {
30+
return dataMasking.getJsonFilters().stream()
31+
.collect(Collectors.toMap(
32+
JsonMaskingFilter::getTopic,
33+
JsonMaskingFilter::getKeys,
34+
(a, b) -> a,
35+
HashMap::new
36+
));
2537
}
2638

2739
public Record maskRecord(Record record) {
2840
try {
29-
if(isJson(record)) {
30-
return jsonMaskingFilters
31-
.stream()
32-
.filter(jsonMaskingFilter -> record.getTopic().getName().equalsIgnoreCase(jsonMaskingFilter.getTopic()))
33-
.findFirst()
34-
.map(filter -> applyMasking(record, filter.getKeys()))
35-
.orElse(record);
41+
if (!isJson(record)) {
42+
return record;
3643
}
44+
return maskJsonRecord(record);
3745
} catch (Exception e) {
38-
LOG.error("Error masking record", e);
46+
LOG.error(ERROR_MESSAGE, e);
47+
return record;
3948
}
40-
return record;
49+
}
50+
51+
private Record maskJsonRecord(Record record) {
52+
String topic = record.getTopic().getName().toLowerCase();
53+
List<String> maskedKeys = topicToKeysMap.get(topic);
54+
return maskedKeys != null ? applyMasking(record, maskedKeys) : record;
4155
}
4256

4357
@SneakyThrows
44-
private Record applyMasking(Record record, List<String> keys) {
45-
JsonObject jsonElement = JsonParser.parseString(record.getValue()).getAsJsonObject();
46-
for(String key : keys) {
47-
maskField(jsonElement, key.split("\\."), 0);
48-
}
49-
record.setValue(jsonElement.toString());
58+
private Record applyMasking(Record record, List<String> maskedKeys) {
59+
JsonObject root = JsonParser.parseString(record.getValue()).getAsJsonObject();
60+
String[][] pathArrays = preProcessPaths(maskedKeys);
61+
maskPaths(root, pathArrays);
62+
record.setValue(root.toString());
5063
return record;
5164
}
5265

53-
private void maskField(JsonObject node, String[] keys, int index) {
54-
if (index == keys.length - 1) {
55-
if (node.has(keys[index])) {
56-
node.addProperty(keys[index], jsonMaskReplacement);
57-
}
66+
private String[][] preProcessPaths(List<String> maskedKeys) {
67+
return maskedKeys.stream()
68+
.map(key -> key.split("\\."))
69+
.toArray(String[][]::new);
70+
}
71+
72+
private void maskPaths(JsonObject root, String[][] pathArrays) {
73+
for (String[] path : pathArrays) {
74+
maskJson(root, path, 0);
75+
}
76+
}
77+
78+
private void maskJson(JsonElement element, String[] path, int index) {
79+
if (index == path.length) return;
80+
81+
String currentKey = path[index];
82+
if (element.isJsonObject()) {
83+
handleJsonObject(element.getAsJsonObject(), path, index, currentKey);
84+
} else if (element.isJsonArray()) {
85+
handleJsonArray(element.getAsJsonArray(), path, index);
86+
}
87+
}
88+
89+
private void handleJsonObject(JsonObject obj, String[] path, int index, String currentKey) {
90+
if (!obj.has(currentKey)) return;
91+
92+
if (index == path.length - 1) {
93+
maskTargetElement(obj, currentKey);
5894
} else {
59-
JsonElement childNode = node.get(keys[index]);
60-
if (childNode != null && childNode.isJsonObject()) {
61-
maskField(childNode.getAsJsonObject(), keys, index + 1);
95+
maskJson(obj.get(currentKey), path, index + 1);
96+
}
97+
}
98+
99+
private void handleJsonArray(JsonArray array, String[] path, int index) {
100+
for (int i = 0; i < array.size(); i++) {
101+
JsonElement arrayElement = array.get(i);
102+
if (arrayElement.isJsonObject()) {
103+
maskJson(arrayElement, path, index);
62104
}
63105
}
64106
}
65-
}
107+
108+
private void maskTargetElement(JsonObject obj, String currentKey) {
109+
JsonElement target = obj.get(currentKey);
110+
if (target.isJsonArray()) {
111+
maskArrayElement(target.getAsJsonArray());
112+
} else {
113+
obj.addProperty(currentKey, jsonMaskReplacement);
114+
}
115+
}
116+
117+
private void maskArrayElement(JsonArray array) {
118+
for (int i = 0; i < array.size(); i++) {
119+
array.set(i, new JsonPrimitive(jsonMaskReplacement));
120+
}
121+
}
122+
}

0 commit comments

Comments
 (0)