Implement radix-trie

Breus · Sep 26, 2024 · ff235e4 · ff235e4
1 parent a0811f0
commit ff235e4
Show file tree

Hide file tree

Showing 7 changed files with 324 additions and 166 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -174,12 +174,17 @@ tasks {
                 "RedundantOverride",
                 "RedundantThrows",
                 "RemoveUnusedImports",
+                "DefaultCharset",
                 "UnnecessarilyFullyQualified",
                 "UnnecessarilyUsedValue",
                 "UnnecessaryBoxedAssignment",
                 "UnnecessaryBoxedVariable",
                 "UnnecessaryFinal",
                 "UnusedException",
+                "UnusedLabel",
+                "UnusedMethod",
+                "UnusedNestedClass",
+                "UnusedVariable",
                 "WildcardImport",
             )
             disable(
@@ -208,4 +213,4 @@ tasks {
     withType<SonarTask> {
         dependsOn(jacocoTestReport)
     }
-}
+}
diff --git a/src/main/java/dev/blaauwendraad/masker/json/KeyMatcher.java b/src/main/java/dev/blaauwendraad/masker/json/KeyMatcher.java
diff --git a/src/main/java/dev/blaauwendraad/masker/json/MaskingState.java b/src/main/java/dev/blaauwendraad/masker/json/MaskingState.java
@@ -56,15 +56,15 @@ final class MaskingState implements ValueMaskerContext {
      * Current JSONPath is represented by a stack of segment references.
      * A stack is implemented with an array of the trie nodes that reference the end of the segment
      */
-    private KeyMatcher.@Nullable TrieNode @Nullable [] currentJsonPath = null;
+    private KeyMatcher.@Nullable StatefulRadixTrieNode @Nullable [] currentJsonPath = null;
     private int currentJsonPathHeadIndex = -1;
     private int currentTokenStartIndex = -1;
 
     public MaskingState(byte[] message, boolean trackJsonPath) {
         this.message = message;
         this.messageLength = message.length;
         if (trackJsonPath) {
-            currentJsonPath = new KeyMatcher.TrieNode[INITIAL_JSONPATH_STACK_CAPACITY];
+            currentJsonPath = new KeyMatcher.StatefulRadixTrieNode[INITIAL_JSONPATH_STACK_CAPACITY];
         }
         this.inputStream = null;
         this.outputStream = null;
@@ -87,7 +87,7 @@ public MaskingState(InputStream inputStream, OutputStream outputStream, boolean
         this.message = new byte[this.bufferSize];
         this.messageLength = 0;
         if (trackJsonPath) {
-            currentJsonPath = new KeyMatcher.TrieNode[INITIAL_JSONPATH_STACK_CAPACITY];
+            currentJsonPath = new KeyMatcher.StatefulRadixTrieNode[INITIAL_JSONPATH_STACK_CAPACITY];
         }
         readNextBuffer();
     }
@@ -232,7 +232,7 @@ boolean jsonPathEnabled() {
      *
      * @param trieNode a node in the trie where the new segment ends.
      */
-    void expandCurrentJsonPath(KeyMatcher.@Nullable TrieNode trieNode) {
+    void expandCurrentJsonPath(KeyMatcher.@Nullable StatefulRadixTrieNode trieNode) {
         if (currentJsonPath != null) {
             currentJsonPath[++currentJsonPathHeadIndex] = trieNode;
             if (currentJsonPathHeadIndex == currentJsonPath.length - 1) {
@@ -254,7 +254,7 @@ void backtrackCurrentJsonPath() {
     /**
      * Returns the TrieNode that references the end of the latest segment in the current jsonpath
      */
-    public KeyMatcher.@Nullable TrieNode getCurrentJsonPathNode() {
+    public KeyMatcher.@Nullable StatefulRadixTrieNode getCurrentJsonPathNode() {
         if (currentJsonPath != null && currentJsonPathHeadIndex != -1) {
             return currentJsonPath[currentJsonPathHeadIndex];
         } else {

diff --git a/src/test/java/dev/blaauwendraad/masker/json/InstanceCreationMemoryUsageTest.java b/src/test/java/dev/blaauwendraad/masker/json/InstanceCreationMemoryUsageTest.java
@@ -1623,7 +1623,7 @@ void realWorldObfuscatedKeysInstanceCreation() {
 
         long memoryBeforeInstanceCreationKb = getCurrentRetainedMemory();
 
-        long memoryLimitKb = 2_000;
+        long memoryLimitKb = 700;
         long memoryConsumedKb = bytesToKb(memoryBeforeInstanceCreationKb - memoryBeforeInstanceCreation);
 
         Assertions.assertThat(memoryConsumedKb)

diff --git a/src/test/java/dev/blaauwendraad/masker/json/JsonMaskerTestUtil.java b/src/test/java/dev/blaauwendraad/masker/json/JsonMaskerTestUtil.java
@@ -39,8 +39,8 @@ public static List<JsonMaskerTestInstance> getJsonMaskerTestInstancesFromFile(St
                 applyConfig(jsonMaskingConfig, builder);
             }
             JsonMaskingConfig maskingConfig = builder.build();
-            var input = jsonNode.get("input").toString();
-            var expectedOutput = jsonNode.get("expectedOutput").toString();
+            var input = jsonNode.get("input").toPrettyString();
+            var expectedOutput = jsonNode.get("expectedOutput").toPrettyString();
             testInstances.add(new JsonMaskerTestInstance(input, expectedOutput, new KeyContainsMasker(maskingConfig)));
         }
         return testInstances;

diff --git a/src/test/java/dev/blaauwendraad/masker/json/KeyMatcherTest.java b/src/test/java/dev/blaauwendraad/masker/json/KeyMatcherTest.java
@@ -111,7 +111,7 @@ void shouldMatchJsonPaths() {
                 """;
         byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
 
-        KeyMatcher.TrieNode node = keyMatcher.getJsonPathRootNode();
+        var node = keyMatcher.getJsonPathRootNode();
         node = keyMatcher.traverseJsonPathSegment(bytes, node, indexOf(bytes, 'a'), 1);
         node = keyMatcher.traverseJsonPathSegment(bytes, node, indexOf(bytes, 'b'), 1);
         assertThat(keyMatcher.getMaskConfigIfMatched(bytes, 0, 0, node)).isNotNull();
@@ -150,7 +150,7 @@ void shouldMatchJsonPathArrays() {
                 """;
         byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
 
-        KeyMatcher.TrieNode node = keyMatcher.getJsonPathRootNode();
+        var node = keyMatcher.getJsonPathRootNode();
         node = keyMatcher.traverseJsonPathSegment(bytes, node, indexOf(bytes, 'a'), 1);
         node = keyMatcher.traverseJsonPathSegment(bytes, node, -1, -1);
         node = keyMatcher.traverseJsonPathSegment(bytes, node, indexOf(bytes, 'b'), 1);
@@ -184,7 +184,7 @@ void shouldNotMatchJsonPathPrefix() {
                 """;
         byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
 
-        KeyMatcher.TrieNode node = keyMatcher.getJsonPathRootNode();
+        var node = keyMatcher.getJsonPathRootNode();
         node = keyMatcher.traverseJsonPathSegment(bytes, node, 2, 4);
         assertThat(keyMatcher.getMaskConfigIfMatched(bytes, 0, -1, node)).isNull();
 
@@ -206,7 +206,7 @@ void shouldReturnMaskingConfigForJsonPathInAllowMode() {
                 """;
         byte[] bytes = json.getBytes(StandardCharsets.UTF_8);
 
-        KeyMatcher.TrieNode node = keyMatcher.getJsonPathRootNode();
+        var node = keyMatcher.getJsonPathRootNode();
         node = keyMatcher.traverseJsonPathSegment(bytes, node, 2, 7);
         assertThat(keyMatcher.getMaskConfigIfMatched(bytes, 0, -1, node)).isNull();
 
@@ -255,4 +255,30 @@ private int indexOf(byte[] bytes, char c) {
         }
         return found;
     }
+
+    @Test
+    void printsNicely() {
+        JsonMaskingConfig config = JsonMaskingConfig.builder()
+                .allowKeys("romane", "romanus", "romulus", "rubens", "ruber", "rubicon", "rubicondus")
+                .build();
+        KeyMatcher keyMatcher = new KeyMatcher(config);
+        assertThat(keyMatcher.printTree())
+                .isEqualTo("""
+                           r -> om -> an -> e
+                                         -> us
+                                   -> ulus
+                             -> ub -> e -> ns
+                                        -> r
+                                   -> icon -> dus
+                           """);
+    }
+
+    @Test
+    void printsEmpty() {
+        JsonMaskingConfig config = JsonMaskingConfig.builder()
+                .allowKeys()
+                .build();
+        KeyMatcher keyMatcher = new KeyMatcher(config);
+        assertThat(keyMatcher.printTree()).isEqualTo("\n");
+    }
 }
diff --git a/src/test/java/dev/blaauwendraad/masker/json/MaskingStateTest.java b/src/test/java/dev/blaauwendraad/masker/json/MaskingStateTest.java
@@ -47,7 +47,7 @@ void shouldReturnStringRepresentationForDebugging() {
     void jsonPathExceedsCapacity() {
         MaskingState maskingState = new MaskingState("[]".getBytes(StandardCharsets.UTF_8), true);
         for (int i = 0; i < 101; i++) {
-            maskingState.expandCurrentJsonPath(KeyMatcher.transform(new KeyMatcher.PreInitTrieNode()));
+            maskingState.expandCurrentJsonPath(new KeyMatcher.StatefulRadixTrieNode(KeyMatcher.compress(new KeyMatcher.PreInitTrieNode())));
         }
         Assertions.assertThat(maskingState.getCurrentJsonPathNode()).isNotNull();
     }
@@ -91,4 +91,4 @@ void shouldUseCorrectOffsetWhenThrowingValueMaskerError() {
                 .hasMessage("Didn't like the value at index 3 at index 19");
     }
 
-}
+}