From c856635ac580217638fd328eeb9468c665430ca7 Mon Sep 17 00:00:00 2001
From: Andrew Steinborn <git@steinborn.me>
Date: Sun, 31 May 2020 08:14:33 -0400
Subject: [PATCH] Try two, with improved state maintenance and some bugs fixed

---
 .../netty/MinecraftVarintFrameDecoder.java    | 82 +++++++++++++------
 1 file changed, 59 insertions(+), 23 deletions(-)

diff --git a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintFrameDecoder.java b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintFrameDecoder.java
index a4da29a76..6cf258e6b 100644
--- a/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintFrameDecoder.java
+++ b/proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintFrameDecoder.java
@@ -1,15 +1,17 @@
 package com.velocitypowered.proxy.protocol.netty;
 
-import com.velocitypowered.proxy.protocol.ProtocolUtils;
 import com.velocitypowered.proxy.util.except.QuietException;
 import io.netty.buffer.ByteBuf;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
+import io.netty.util.ByteProcessor;
 import java.util.List;
 
 public class MinecraftVarintFrameDecoder extends ByteToMessageDecoder {
 
   private static final QuietException BAD_LENGTH_CACHED = new QuietException("Bad packet length");
+  private static final QuietException VARINT_BIG_CACHED = new QuietException("VarInt too big");
+  private final VarintByteDecoder reader = new VarintByteDecoder();
 
   @Override
   protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) {
@@ -18,31 +20,65 @@ public class MinecraftVarintFrameDecoder extends ByteToMessageDecoder {
       return;
     }
 
-    read_lens: while (in.isReadable()) {
-      int origReaderIndex = in.readerIndex();
-      for (int i = 0; i < 3; i++) {
-        if (!in.isReadable()) {
-          in.readerIndex(origReaderIndex);
-          return;
-        }
+    while (in.isReadable()) {
+      reader.reset();
 
-        byte read = in.readByte();
-        if (read >= 0) {
-          // Make sure reader index of length buffer is returned to the beginning
-          in.readerIndex(origReaderIndex);
-          int packetLength = ProtocolUtils.readVarInt(in);
-
-          if (in.readableBytes() >= packetLength) {
-            out.add(in.readBytes(packetLength));
-            continue read_lens;
-          } else {
-            in.readerIndex(origReaderIndex);
-            return;
-          }
-        }
+      int varintEnd = in.forEachByte(reader);
+      if (varintEnd == -1) {
+        // We tried to go beyond the end of the buffer. This is probably a good sign that the
+        // buffer was too short to hold a proper varint.
+        return;
       }
 
-      throw BAD_LENGTH_CACHED;
+      if (reader.result == DecodeResult.SUCCESS) {
+        if (reader.readVarint < 0) {
+          throw BAD_LENGTH_CACHED;
+        }
+
+        int minimumRead = reader.bytesRead + reader.readVarint;
+        if (in.isReadable(minimumRead)) {
+          out.add(in.retainedSlice(varintEnd + 1, reader.readVarint));
+          in.skipBytes(minimumRead);
+        } else {
+          return;
+        }
+      } else if (reader.result == DecodeResult.TOO_BIG) {
+        throw VARINT_BIG_CACHED;
+      } else if (reader.result == DecodeResult.TOO_SHORT) {
+        // No-op: we couldn't get a useful result.
+      }
     }
   }
+
+  private static class VarintByteDecoder implements ByteProcessor {
+    private int readVarint;
+    private int bytesRead;
+    private DecodeResult result = DecodeResult.TOO_SHORT;
+
+    @Override
+    public boolean process(byte k) {
+      readVarint |= (k & 0x7F) << bytesRead++ * 7;
+      if (bytesRead > 3) {
+        result = DecodeResult.TOO_BIG;
+        return false;
+      }
+      if ((k & 0x80) != 128) {
+        result = DecodeResult.SUCCESS;
+        return false;
+      }
+      return true;
+    }
+
+    void reset() {
+      readVarint = 0;
+      bytesRead = 0;
+      result = DecodeResult.TOO_SHORT;
+    }
+  }
+
+  private enum DecodeResult {
+    SUCCESS,
+    TOO_SHORT,
+    TOO_BIG
+  }
 }