Fix geyser connection timeout caused by translation expansion

core/src/main/java/org/geysermc/geyser/translator/text/MessageTranslator.java

@@ -29,6 +29,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.function.Consumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import net.kyori.adventure.text.Component;
@@ -79,6 +80,8 @@ public class MessageTranslator {
     // Reset character
     private static final String RESET = BASE + "r";
     private static final Pattern LOCALIZATION_PATTERN = Pattern.compile("%(?:(\\d+)\\$)?s");
+    private static final int MAX_TRANSLATION_AMPLIFICATION = 4_096;
+    private static final ThreadLocal<long[]> TRANSLATION_AMPLIFICATION = new ThreadLocal<>();
 
     static {
         GSON_SERIALIZER = DefaultComponentSerializer.get()
@@ -116,6 +119,19 @@ public class MessageTranslator {
                 final String translated = translatable.key();
                 final Matcher matcher = LOCALIZATION_PATTERN.matcher(translated);
                 final List<TranslationArgument> args = translatable.arguments();
+                final int[] occurrences = new int[args.size()];
+                int occurrenceArgPosition = 0;
+                while (matcher.find()) {
+                    try {
+                        final String argIdx = matcher.group(1);
+                        final int idx = argIdx != null ? Integer.parseInt(argIdx) - 1 : occurrenceArgPosition++;
+                        if (idx >= 0 && idx < occurrences.length) {
+                            occurrences[idx]++;
+                        }
+                    } catch (final NumberFormatException ignored) {
+                    }
+                }
+                matcher.reset();
                 int argPosition = 0;
                 int lastIdx = 0;
                 while (matcher.find()) {
@@ -130,16 +146,16 @@ public class MessageTranslator {
                     if (argIdx != null) {
                         try {
                             final int idx = Integer.parseInt(argIdx) - 1;
-                            if (idx < args.size()) {
-                                consumer.accept(args.get(idx).asComponent());
+                            if (idx >= 0 && idx < args.size()) {
+                                acceptTranslationArgument(args.get(idx).asComponent(), occurrences[idx], consumer);
                             }
                         } catch (final NumberFormatException ex) {
                             // ignore, drop the format placeholder
                         }
                     } else {
                         final int idx = argPosition++;
                         if (idx < args.size()) {
-                            consumer.accept(args.get(idx).asComponent());
+                            acceptTranslationArgument(args.get(idx).asComponent(), occurrences[idx], consumer);
                         }
                     }
                 }
@@ -221,7 +237,17 @@ private static String convertMessage(Component message, String locale, boolean a
             // Translate any components that require it
             message = RENDERER.render(message, locale);
 
-            String legacy = BEDROCK_SERIALIZER.serialize(message);
+            long[] translationAmplification = {1, 0};
+            TRANSLATION_AMPLIFICATION.set(translationAmplification);
+            String legacy;
+            try {
+                legacy = BEDROCK_SERIALIZER.serialize(message);
+            } finally {
+                TRANSLATION_AMPLIFICATION.remove();
+            }
+            if (translationAmplification[1] != 0) {
+                return "";
+            }
             int legacyLength = legacy.length();
 
             // We need to allocate at least the length of the original message, it can only grow
@@ -308,6 +334,28 @@ private static int setFormattingFlag(int flags, char format) {
         return 1 << index;
     }
 
+    private static void acceptTranslationArgument(Component argument, int occurrences, Consumer<Component> consumer) {
+        long[] amplification = TRANSLATION_AMPLIFICATION.get();
+        if (amplification == null) {
+            consumer.accept(argument);
+            return;
+        }
+
+        long previous = amplification[0];
+        long current = Math.min(MAX_TRANSLATION_AMPLIFICATION + 1L, previous * occurrences);
+        if (current > MAX_TRANSLATION_AMPLIFICATION) {
+            amplification[1] = 1;
+            return;
+        }
+
+        amplification[0] = current;
+        try {
+            consumer.accept(argument);
+        } finally {
+            amplification[0] = previous;
+        }
+    }
+
     private static void applyFormattingFlags(int flags, StringBuilder builder) {
         int colorCount = BEDROCK_COLORS.length();
 

core/src/test/java/org/geysermc/geyser/network/translators/chat/MessageTranslatorTest.java

@@ -25,8 +25,11 @@
 
 package org.geysermc.geyser.network.translators.chat;
 
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import net.kyori.adventure.text.Component;
 import org.geysermc.geyser.translator.text.MessageTranslator;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -121,6 +124,31 @@ public void convertMessage() {
         }
     }
 
+    @Test
+    public void rejectExcessiveTranslationExpansion() {
+        Component message = Component.text("x");
+        for (int i = 0; i < 5; i++) {
+            message = Component.translatable("%1$s".repeat(10), message);
+        }
+        Assertions.assertEquals("", MessageTranslator.convertMessage(message, "en_US"));
+    }
+
+    @Test
+    public void allowWideTranslation() {
+        List<Component> arguments = Collections.nCopies(5_000, Component.text("x"));
+        Component message = Component.translatable("%s".repeat(arguments.size()), arguments);
+
+        Assertions.assertEquals("x".repeat(arguments.size()), MessageTranslator.convertMessageRaw(message, "en_US"));
+    }
+
+    @Test
+    public void allowDetailedTranslation() {
+        String name = "x".repeat(70_000);
+        Component message = Component.translatable("%1$s", Component.text(name));
+
+        Assertions.assertEquals(name, MessageTranslator.convertMessageRaw(message, "en_US"));
+    }
+
     @Test
     public void convertMessageLenient() {
         Assertions.assertEquals("\n\n\n\n", MessageTranslator.convertMessageLenient("\n\n\n\n"), "All newline message is not handled properly");