Skip to content

Commit

Permalink
Split json path state into separate class
Browse files Browse the repository at this point in the history
  • Loading branch information
gavlyukovskiy committed Dec 16, 2024
1 parent 9d0d51c commit db1f605
Show file tree
Hide file tree
Showing 11 changed files with 509 additions and 384 deletions.
4 changes: 4 additions & 0 deletions build.gradle.kts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import net.ltgt.gradle.errorprone.errorprone
import org.gradle.nativeplatform.platform.internal.DefaultNativePlatform
import org.sonarqube.gradle.SonarTask

plugins {
Expand Down Expand Up @@ -187,6 +188,9 @@ tasks {
"UnusedVariable",
"WildcardImport",
)
if (DefaultNativePlatform.getCurrentOperatingSystem().isWindows) {
disable("MisleadingEscapedSpace") // good stuff
}
disable(
"StringCaseLocaleUsage",
"MissingSummary",
Expand Down
103 changes: 103 additions & 0 deletions src/main/java/dev/blaauwendraad/masker/json/JsonPathState.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
package dev.blaauwendraad.masker.json;

import org.jspecify.annotations.Nullable;

import java.util.ArrayDeque;

class JsonPathState {
private static final KeyMatcher.StatefulRadixTrieNode NULL_NODE = new KeyMatcher.StatefulRadixTrieNode(new KeyMatcher.RadixTrieNode(new byte[0], new byte[0]), 0);

private final KeyMatcher keyMatcher;
private final ArrayDeque<KeyMatcher.StatefulRadixTrieNode> jsonPathSegments = new ArrayDeque<>();

JsonPathState(KeyMatcher keyMatcher) {
this.keyMatcher = keyMatcher;
var root = keyMatcher.getRootNode();
if (!root.performChildLookup((byte) '$')) {
throw new IllegalStateException("JSONPath root node is null");
}
this.jsonPathSegments.push(new KeyMatcher.StatefulRadixTrieNode(root));
root.reset();
}

/**
* Expands current JSONPath with an array segment.
*/
void pushArraySegment() {
jsonPathSegments.push(getWildcardNodeOrNullNode());
}

/**
* Expands current JSONPath with a value segment.
*/
void pushKeyValueSegment(byte[] bytes, int keyOffset, int keyLength) {
jsonPathSegments.push(getKeyValueNodeOrNullNode(bytes, keyOffset, keyLength));
}

/**
* Backtracks current JSONPath to the previous segment.
*/
void backtrack() {
jsonPathSegments.pop();
}

/**
* Traverse the trie node when entering an array. In order to match the array it has to be a wildcard.
* <p>For example:
* For a JSON like this {@code { "holder": [ { "maskMe": "secret" } } } the matching JSONPath has to be
* {@code '$.holder.*.maskMe'}, so that entering the array requires a wildcard node.
*/
private KeyMatcher.StatefulRadixTrieNode getWildcardNodeOrNullNode() {
var current = currentNode();
if (current == null) {
return NULL_NODE;
}
try {
if (!current.performChildLookup((byte) '.')) {
return NULL_NODE;
}
if (current.isJsonPathWildcard()) {
current.performChildLookup((byte) '*');
return new KeyMatcher.StatefulRadixTrieNode(current);
}
return NULL_NODE;
} finally {
current.reset();
}
}

/**
* Traverse the trie node when entering a key-value. The matching can be done for the matching key, or through a wildcard ('*') JSONPath.
*/
private KeyMatcher.StatefulRadixTrieNode getKeyValueNodeOrNullNode(byte[] bytes, int keyOffset, int keyLength) {
var current = currentNode();
if (current == null) {
return NULL_NODE;
}
try {
if (!current.performChildLookup((byte) '.')) {
return NULL_NODE;
}
if (current.isJsonPathWildcard()) {
current.performChildLookup((byte) '*');
return new KeyMatcher.StatefulRadixTrieNode(current);
} else {
var child = keyMatcher.traverseFrom(current, bytes, keyOffset, keyLength);
if (child != null) {
return new KeyMatcher.StatefulRadixTrieNode(child);
}
}
return NULL_NODE;
} finally {
current.reset();
}
}

KeyMatcher.@Nullable StatefulRadixTrieNode currentNode() {
var peek = jsonPathSegments.peek();
if (peek == NULL_NODE) {
return null;
}
return peek;
}
}
63 changes: 40 additions & 23 deletions src/main/java/dev/blaauwendraad/masker/json/KeyContainsMasker.java
Original file line number Diff line number Diff line change
Expand Up @@ -66,14 +66,18 @@ public void mask(InputStream inputStream, OutputStream outputStream) {
private void mask(MaskingState maskingState) {
try {
KeyMaskingConfig keyMaskingConfig = maskingConfig.isInAllowMode() ? maskingConfig.getDefaultConfig() : null;
if (maskingState.jsonPathEnabled()) {
maskingState.expandCurrentJsonPath(keyMatcher.getJsonPathRootNode());
keyMaskingConfig = keyMatcher.getMaskConfigIfMatched(maskingState.getMessage(), -1, -1, maskingState.getCurrentJsonPathNode());

JsonPathState jsonPathState;
if (!maskingConfig.getTargetJsonPaths().isEmpty()) {
jsonPathState = new JsonPathState(keyMatcher);
keyMaskingConfig = keyMatcher.getMaskConfigIfMatched(maskingState.getMessage(), -1, -1, jsonPathState.currentNode());
} else {
jsonPathState = null;
}

while (!maskingState.endOfJson()) {
stepOverWhitespaceCharacters(maskingState);
if (!visitValue(maskingState, keyMaskingConfig)) {
if (!visitValue(maskingState, jsonPathState, keyMaskingConfig)) {
maskingState.next();
}
}
Expand All @@ -86,19 +90,19 @@ private void mask(MaskingState maskingState) {
* Entrypoint of visiting any value (object, array or primitive) in the JSON.
*
* @param maskingState the current masking state
* @param jsonPathState the current {@link JsonPathState}
* @param keyMaskingConfig if not null it means that the current value is being masked otherwise the value is not
* being masked
*
* @return whether a value was found, if returned false the calling code must advance to avoid infinite loops
*/
private boolean visitValue(MaskingState maskingState, @Nullable KeyMaskingConfig keyMaskingConfig) {
private boolean visitValue(MaskingState maskingState, @Nullable JsonPathState jsonPathState, @Nullable KeyMaskingConfig keyMaskingConfig) {
if (maskingState.endOfJson()) {
return true;
}
// using switch-case over 'if'-statements to improve performance by ~20% (measured in benchmarks)
switch (maskingState.byteAtCurrentIndex()) {
case '[' -> visitArray(maskingState, keyMaskingConfig);
case '{' -> visitObject(maskingState, keyMaskingConfig);
case '[' -> visitArray(maskingState, jsonPathState, keyMaskingConfig);
case '{' -> visitObject(maskingState, jsonPathState, keyMaskingConfig);
case '-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9' -> {
if (keyMaskingConfig != null) {
maskNumber(maskingState, keyMaskingConfig);
Expand Down Expand Up @@ -136,23 +140,26 @@ private boolean visitValue(MaskingState maskingState, @Nullable KeyMaskingConfig
}

/**
* Visits an array of unknown values (or empty) and invokes {@link #visitValue(MaskingState, KeyMaskingConfig)} on
* Visits an array of unknown values (or empty) and invokes {@link #visitValue(MaskingState, JsonPathState, KeyMaskingConfig)} on
* each element while propagating the {@link KeyMaskingConfig}.
*
* @param maskingState the current {@link MaskingState}
* @param jsonPathState the current {@link JsonPathState}
* @param keyMaskingConfig if not null it means that the current value is being masked according to the
* {@link KeyMaskingConfig}. Otherwise, the value is not masked
*/
private void visitArray(MaskingState maskingState, @Nullable KeyMaskingConfig keyMaskingConfig) {
maskingState.expandCurrentJsonPath(keyMatcher.traverseJsonPathSegment(maskingState.getMessage(), maskingState.getCurrentJsonPathNode(), -1, -1));
private void visitArray(MaskingState maskingState, @Nullable JsonPathState jsonPathState, @Nullable KeyMaskingConfig keyMaskingConfig) {
if (jsonPathState != null) {
jsonPathState.pushArraySegment();
}
while (maskingState.next()) {
stepOverWhitespaceCharacters(maskingState);
// check if we're in an empty array
if (maskingState.byteAtCurrentIndex() == ']') {
break;
}

visitValue(maskingState, keyMaskingConfig);
visitValue(maskingState, jsonPathState, keyMaskingConfig);

stepOverWhitespaceCharacters(maskingState);
// check if we're at the end of a (non-empty) array
Expand All @@ -161,22 +168,25 @@ private void visitArray(MaskingState maskingState, @Nullable KeyMaskingConfig ke
}
}
maskingState.next(); // step over array closing square bracket
maskingState.backtrackCurrentJsonPath();
if (jsonPathState != null) {
jsonPathState.backtrack();
}
}

/**
* Visits an object, iterates over the keys and checks whether key needs to be masked (if
* {@link JsonMaskingConfig.TargetKeyMode#MASK}) or allowed (if {@link JsonMaskingConfig.TargetKeyMode#ALLOW}). For
* each value, invokes {@link #visitValue(MaskingState, KeyMaskingConfig)} with a non-null {@link KeyMaskingConfig}
* each value, invokes {@link #visitValue(MaskingState, JsonPathState, KeyMaskingConfig)} with a non-null {@link KeyMaskingConfig}
* (when key needs to be masked) or {@code null} (when key is allowed). Whenever 'parentKeyMaskingConfig' is
* supplied, it means that the object with all its keys is being masked. The only situation when the individual
* values do not need to be masked is when the key is explicitly allowed (in allow mode).
*
* @param maskingState the current {@link MaskingState}
* @param jsonPathState the current {@link JsonPathState}
* @param parentKeyMaskingConfig if not null it means that the current value is being masked according to the
* {@link KeyMaskingConfig}. Otherwise, the value is not being masked
*/
private void visitObject(MaskingState maskingState, @Nullable KeyMaskingConfig parentKeyMaskingConfig) {
private void visitObject(MaskingState maskingState, @Nullable JsonPathState jsonPathState, @Nullable KeyMaskingConfig parentKeyMaskingConfig) {
while (maskingState.next()) {
stepOverWhitespaceCharacters(maskingState);
// check if we're in an empty object
Expand All @@ -188,12 +198,17 @@ private void visitObject(MaskingState maskingState, @Nullable KeyMaskingConfig p

stepOverStringValue(maskingState);

int keyStartIndex = maskingState.getCurrentTokenStartIndex();
int afterClosingQuoteIndex = maskingState.currentIndex();
int keyLength = afterClosingQuoteIndex - keyStartIndex - 2; // minus the opening and closing quotes
maskingState.expandCurrentJsonPath(keyMatcher.traverseJsonPathSegment(maskingState.getMessage(), maskingState.getCurrentJsonPathNode(), keyStartIndex + 1, keyLength));
KeyMaskingConfig keyMaskingConfig = keyMatcher.getMaskConfigIfMatched(maskingState.getMessage(), keyStartIndex + 1, // plus one for the opening quote
keyLength, maskingState.getCurrentJsonPathNode());
int keyStartIndex = maskingState.getCurrentTokenStartIndex() + 1; // plus the opening quote
int keyLength = maskingState.currentIndex() - keyStartIndex - 1; // minus the closing quote
KeyMaskingConfig keyMaskingConfig;
if (jsonPathState != null) {
jsonPathState.pushKeyValueSegment(maskingState.getMessage(), keyStartIndex, keyLength);
keyMaskingConfig = keyMatcher.getMaskConfigIfMatched(maskingState.getMessage(), keyStartIndex, keyLength, jsonPathState.currentNode());
} else {
keyMaskingConfig = keyMatcher.getMaskConfigIfMatched(maskingState.getMessage(), keyStartIndex, keyLength, null);
}


maskingState.clearTokenStartIndex();
stepOverWhitespaceCharacters(maskingState);
// step over the colon ':'
Expand All @@ -214,9 +229,11 @@ private void visitObject(MaskingState maskingState, @Nullable KeyMaskingConfig p
if (parentKeyMaskingConfig != null && (keyMaskingConfig == null || keyMaskingConfig == maskingConfig.getDefaultConfig())) {
keyMaskingConfig = parentKeyMaskingConfig;
}
visitValue(maskingState, keyMaskingConfig);
visitValue(maskingState, jsonPathState, keyMaskingConfig);
}
if (jsonPathState != null) {
jsonPathState.backtrack();
}
maskingState.backtrackCurrentJsonPath();

stepOverWhitespaceCharacters(maskingState);
// check if we're at the end of a (non-empty) object
Expand Down
Loading

0 comments on commit db1f605

Please sign in to comment.