Reject trie nodes that exceed the branch-count limit

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;
 

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);