Support negative numbers in writeVLong (#22314)

nik9000 · nik9000 · commit 9276290c0fcd · 2016-12-22T14:14:41.000-05:00
We don't *want* to use negative numbers with `writeVLong`
so throw an exception when we try. On the other
hand unforeseen bugs might cause us to write negative numbers (some versions of Elasticsearch don't have the exception, only an assertion)
so this fixes `readVLong` so that instead of reading a wrong
value and corrupting the stream it reads the negative value.
diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamInput.java
@@ -214,9 +214,8 @@ public long readLong() throws IOException {
     }
 
     /**
-     * Reads a long stored in variable-length format.  Reads between one and
-     * nine bytes.  Smaller values take fewer bytes.  Negative numbers are not
-     * supported.
+     * Reads a long stored in variable-length format. Reads between one and ten bytes. Smaller values take fewer bytes. Negative numbers
+     * are encoded in ten bytes so prefer {@link #readLong()} or {@link #readZLong()} for negative numbers.
      */
     public long readVLong() throws IOException {
         byte b = readByte();
@@ -260,8 +259,16 @@ public long readVLong() throws IOException {
             return i;
         }
         b = readByte();
-        assert (b & 0x80) == 0;
-        return i | ((b & 0x7FL) << 56);
+        i |= ((b & 0x7FL) << 56);
+        if ((b & 0x80) == 0) {
+            return i;
+        }
+        b = readByte();
+        if (b != 0 && b != 1) {
+            throw new IOException("Invalid vlong (" + Integer.toHexString(b) + " << 63) | " + Long.toHexString(i));
+        }
+        i |= ((long) b) << 63;
+        return i;
     }
 
     public long readZLong() throws IOException {
diff --git a/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java b/core/src/main/java/org/elasticsearch/common/io/stream/StreamOutput.java
@@ -210,12 +210,22 @@ public void writeLong(long i) throws IOException {
     }
 
     /**
-     * Writes a non-negative long in a variable-length format.
-     * Writes between one and nine bytes. Smaller values take fewer bytes.
-     * Negative numbers are not supported.
+     * Writes a non-negative long in a variable-length format. Writes between one and ten bytes. Smaller values take fewer bytes. Negative
+     * numbers use ten bytes and trip assertions (if running in tests) so prefer {@link #writeLong(long)} or {@link #writeZLong(long)} for
+     * negative numbers.
      */
     public void writeVLong(long i) throws IOException {
-        assert i >= 0;
+        if (i < 0) {
+            throw new IllegalStateException("Negative longs unsupported, use writeLong or writeZLong for negative numbers [" + i + "]");
+        }
+        writeVLongNoCheck(i);
+    }
+
+    /**
+     * Writes a long in a variable-length format without first checking if it is negative. Package private for testing. Use
+     * {@link #writeVLong(long)} instead.
+     */
+    void writeVLongNoCheck(long i) throws IOException {
         while ((i & ~0x7F) != 0) {
             writeByte((byte) ((i & 0x7f) | 0x80));
             i >>>= 7;
diff --git a/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java b/core/src/test/java/org/elasticsearch/common/io/stream/BytesStreamsTests.java
@@ -21,7 +21,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.Constants;
-import org.apache.lucene.util.UnicodeUtil;
+import org.elasticsearch.Version;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.geo.GeoPoint;
@@ -33,6 +33,7 @@
 import java.io.EOFException;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
@@ -778,4 +779,37 @@ public void testReadNegativeArraySize() throws IOException {
             }
         }
     }
+
+    public void testVInt() throws IOException {
+        final int value = randomInt();
+        BytesStreamOutput output = new BytesStreamOutput();
+        output.writeVInt(value);
+        StreamInput input = output.bytes().streamInput();
+        assertEquals(value, input.readVInt());
+    }
+
+    public void testVLong() throws IOException {
+        final long value = randomLong();
+        {
+            // Read works for positive and negative numbers
+            BytesStreamOutput output = new BytesStreamOutput();
+            output.writeVLongNoCheck(value); // Use NoCheck variant so we can write negative numbers
+            StreamInput input = output.bytes().streamInput();
+            assertEquals(value, input.readVLong());
+        }
+        if (value < 0) {
+            // Write doesn't work for negative numbers
+            BytesStreamOutput output = new BytesStreamOutput();
+            Exception e = expectThrows(IllegalStateException.class, () -> output.writeVLong(value));
+            assertEquals("Negative longs unsupported, use writeLong or writeZLong for negative numbers [" + value + "]", e.getMessage());
+        }
+
+        assertTrue("If we're not compatible with 5.1.1 we can drop the assertion below",
+                Version.CURRENT.minimumCompatibilityVersion().onOrBefore(Version.V_5_1_1_UNRELEASED));
+        /* Read -1 as serialized by a version of Elasticsearch that supported writing negative numbers with writeVLong. Note that this
+         * should be the same test as the first case (when value is negative) but we've kept some bytes so no matter what we do to
+         * writeVLong in the future we can be sure we can read bytes as written by Elasticsearch before 5.1.2 */
+        StreamInput in = new BytesArray(Base64.getDecoder().decode("////////////AQAAAAAAAA==")).streamInput();
+        assertEquals(-1, in.readVLong());
+    }
 }
