From 5cd012736e4680bbb25928ee7fbcea4859878fff Mon Sep 17 00:00:00 2001
From: Alex Sorokoumov <aleksandr.sorokoumov@gmail.com>
Date: Tue, 1 Nov 2022 11:52:31 +0100
Subject: [PATCH] VIntCoding handles BB with less than 8 bytes

Without this fix VIntCoding#writeUnsignedVInt(long, ByteBuffer) throws
a runtime error if the BB size is less than 8 bytes. This method also
silently does not write to BB if it has less than 8 bytes left.

patch by Alex Sorokoumov; reviewed by Benjamin Lerer for CASSANDRA-15215
---
 .../cassandra/utils/vint/VIntCoding.java      | 37 ++++++++++++++++-
 .../cassandra/utils/vint/VIntCodingTest.java  | 40 +++++++++++++++++++
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/src/java/org/apache/cassandra/utils/vint/VIntCoding.java b/src/java/org/apache/cassandra/utils/vint/VIntCoding.java
index 60e8c0de83..8543e6f127 100644
--- a/src/java/org/apache/cassandra/utils/vint/VIntCoding.java
+++ b/src/java/org/apache/cassandra/utils/vint/VIntCoding.java
@@ -50,6 +50,7 @@ import java.io.DataInput;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 
+import io.netty.util.concurrent.FastThreadLocal;
 import net.nicoulaj.compilecommand.annotations.Inline;
 import org.apache.cassandra.io.util.DataInputPlus;
 import org.apache.cassandra.io.util.DataOutputPlus;
@@ -60,6 +61,16 @@ import org.apache.cassandra.io.util.DataOutputPlus;
  */
 public class VIntCoding
 {
+
+    protected static final FastThreadLocal<byte[]> encodingBuffer = new FastThreadLocal<byte[]>()
+    {
+        @Override
+        public byte[] initialValue()
+        {
+            return new byte[9];
+        }
+    };
+
     public static final int MAX_SIZE = 10;
 
     public static long readUnsignedVInt(DataInput input) throws IOException
@@ -214,7 +225,7 @@ public class VIntCoding
         {
             int limit = output.limit();
             int pos = output.position();
-            if (limit - pos >= size)
+            if (limit - pos >= 8)
             {
                 int shift = (8 - size) << 3;
                 int extraBytes = size - 1;
@@ -223,6 +234,10 @@ public class VIntCoding
                 output.putLong(pos, register);
                 output.position(pos + size);
             }
+            else
+            {
+                output.put(VIntCoding.encodeUnsignedVInt(value, size), 0, size);
+            }
         }
         else if (size == 9)
         {
@@ -247,6 +262,26 @@ public class VIntCoding
         writeUnsignedVInt(encodeZigZag64(value), output);
     }
 
+    /**
+     * @return a TEMPORARY THREAD LOCAL BUFFER containing the encoded bytes of the value
+     * This byte[] must be discarded by the caller immediately, and synchronously
+     */
+    @Inline
+    private static byte[] encodeUnsignedVInt(long value, int size)
+    {
+        byte[] encodingSpace = encodingBuffer.get();
+
+        int extraBytes = size - 1;
+        for (int i = extraBytes ; i >= 0; --i)
+        {
+            encodingSpace[i] = (byte) value;
+            value >>= 8;
+        }
+        encodingSpace[0] |= VIntCoding.encodeExtraBytesToRead(extraBytes);
+
+        return encodingSpace;
+    }
+
     /**
      * Decode a ZigZag-encoded 64-bit value.  ZigZag encodes signed integers
      * into values that can be efficiently encoded with varint.  (Otherwise,
diff --git a/test/unit/org/apache/cassandra/utils/vint/VIntCodingTest.java b/test/unit/org/apache/cassandra/utils/vint/VIntCodingTest.java
index bb85b94337..15f9cdc9d3 100644
--- a/test/unit/org/apache/cassandra/utils/vint/VIntCodingTest.java
+++ b/test/unit/org/apache/cassandra/utils/vint/VIntCodingTest.java
@@ -22,6 +22,7 @@ import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.BufferOverflowException;
 import java.nio.ByteBuffer;
 
 import org.apache.cassandra.io.util.DataInputPlus;
@@ -32,6 +33,8 @@ import org.junit.Test;
 
 import org.junit.Assert;
 
+import static org.junit.Assert.fail;
+
 public class VIntCodingTest
 {
     private static final long[] LONGS = new long[] {53L, 10201L, 1097151L,
@@ -161,4 +164,41 @@ public class VIntCodingTest
             Assert.assertEquals(val, VIntCoding.readUnsignedVInt(new DataInputPlus.DataInputStreamPlus(is)));
         }
     }
+
+    @Test
+    public void testWriteUnsignedVIntBBLessThan8Bytes() throws IOException
+    {
+        long val = 10201L;
+        Assert.assertEquals(2, VIntCoding.computeUnsignedVIntSize(val));
+        ByteBuffer bb = ByteBuffer.allocate(2);
+        VIntCoding.writeUnsignedVInt(val, bb);
+        // read as ByteBuffer
+        Assert.assertEquals(val, VIntCoding.getUnsignedVInt(bb, 0));
+        // read as DataInput
+        InputStream is = new ByteArrayInputStream(bb.array());
+        Assert.assertEquals(val, VIntCoding.readUnsignedVInt(new DataInputPlus.DataInputStreamPlus(is)));
+    }
+
+    @Test
+    public void testWriteUnsignedVIntBBHasLessThan8BytesLeft()
+    {
+        long val = 10201L;
+        Assert.assertEquals(2, VIntCoding.computeUnsignedVIntSize(val));
+        ByteBuffer bb = ByteBuffer.allocate(3);
+        bb.position(1);
+        VIntCoding.writeUnsignedVInt(val, bb);
+        // read as ByteBuffer
+        Assert.assertEquals(val, VIntCoding.getUnsignedVInt(bb, 1));
+    }
+
+    @Test
+    public void testWriteUnsignedVIntBBDoesNotHaveEnoughSpaceOverflows()
+    {
+        ByteBuffer bb = ByteBuffer.allocate(3);
+        try
+        {
+            VIntCoding.writeUnsignedVInt(52057592037927932L, bb);
+            fail();
+        } catch (BufferOverflowException e) {}
+    }
 }