diff --git a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java index 5e3f617..6f802c6 100644 --- a/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java +++ b/utils/src/main/java/datadog/instrument/utils/ClassNameTrie.java @@ -104,6 +104,9 @@ public final class ClassNameTrie { /** Marks a glob in the trie, where a match succeeds even if the key has extra characters. */ private static final char GLOB_MARKER = 0x2000; + /** Maximum number of branches in a single node of the trie. */ + private static final int MAX_NODE_BRANCHES = 32767; + /** Maximum value that can be held in a single node of the trie. */ private static final char MAX_NODE_VALUE = 0x1FFF; @@ -606,6 +609,10 @@ private int insertBranch( int newBranch, int subTrieEnd) { + if (branchCount >= MAX_NODE_BRANCHES) { + throw new IllegalArgumentException("Class-name does not fit in the trie"); + } + int remainingKeyLength = key.length() - keyIndex; int insertedCharacters = 3 /* segment-length, jump, value */ + remainingKeyLength; diff --git a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java index 43e48be..45865a6 100644 --- a/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java +++ b/utils/src/test/java/datadog/instrument/utils/ClassNameTrieTest.java @@ -276,6 +276,24 @@ void overflow() { } } + @Test + void tooManyBranchesPreservesExistingMappings() throws IOException { + ClassNameTrie.Builder builder = new ClassNameTrie.Builder(); + // 32767 siblings after a segment; use non-ASCII characters to avoid '*' and '/'. + IntStream.range(0x8000, 0xFFFF).forEach(c -> builder.put("ab" + (char) c, 1)); + byte[] before = serialize(builder); + + assertThrows(IllegalArgumentException.class, () -> builder.put("ab\uFFFF", 1)); + assertArrayEquals(before, serialize(builder)); + assertEquals(-1, builder.apply("ab\uFFFF")); + + builder.put("ab\u8000", 2); // overwrites do not add a sibling + builder.put("ab\u8000x", 3); // extending a branch remains valid at the limit + assertEquals(3, builder.apply("ab\u8000x")); + IntStream.range(0x8000, 0xFFFF) + .forEach(c -> assertEquals(c == 0x8000 ? 2 : 1, builder.apply("ab" + (char) c))); + } + @Test void tooManyLongJumps() { String padding = padding(7700); @@ -437,15 +455,12 @@ private static ClassNameTrie.Builder builderAtLongJumpLimit(String padding) { @MethodSource("insertionsGrowingAncestorJump") void failedAncestorJumpConversionPreservesTrie(String key) throws IOException { ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(0); - ByteArrayOutputStream before = new ByteArrayOutputStream(); - builder.writeTo(new DataOutputStream(before)); + byte[] before = serialize(builder); // These insertions need no local long jump, but grow an ancestor's short jump past the limit. for (int attempt = 0; attempt < 2; attempt++) { assertThrows(IllegalArgumentException.class, () -> builder.put(key, 2)); - ByteArrayOutputStream after = new ByteArrayOutputStream(); - builder.writeTo(new DataOutputStream(after)); - assertArrayEquals(before.toByteArray(), after.toByteArray()); + assertArrayEquals(before, serialize(builder)); } assertEquals(1, builder.apply("aa" + padding(7675))); assertEquals(1, builder.apply("z")); @@ -484,6 +499,12 @@ private static String padding(int length) { return new String(chars); } + private static byte[] serialize(ClassNameTrie.Builder builder) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + builder.writeTo(new DataOutputStream(out)); + return out.toByteArray(); + } + @Test void fullLongJumpTableAllowsSmallInsertionsWithoutConversion() { ClassNameTrie.Builder builder = builderWithShortJumpNearLimit(0, 10);