Merge pull request #168 from pontusmelke/1.0-bitmask-error

pontusmelke · pontusmelke · commit c84db27281aa · 2016-04-22T16:12:45.000+02:00
Fixed bitmasking error
diff --git a/driver/src/main/java/org/neo4j/driver/internal/connector/socket/BufferingChunkedInput.java b/driver/src/main/java/org/neo4j/driver/internal/connector/socket/BufferingChunkedInput.java
@@ -88,6 +88,15 @@ public BufferingChunkedInput( ReadableByteChannel channel, int bufferCapacity )
         this.state = State.AWAITING_CHUNK;
     }
 
+    /*
+     * Use only in tests
+     */
+    int remainingChunkSize()
+    {
+        return remainingChunkSize;
+    }
+
+
     /**
      * Internal state machine used for reading data from the channel into the buffer.
      */
@@ -118,7 +127,7 @@ else if ( ctx.buffer.remaining() >= 2 )
                         {
                             //only 1 byte in buffer, read that and continue
                             //to read header
-                            byte partialChunkSize = ctx.buffer.get();
+                            int partialChunkSize = getUnsignedByteFromBuffer( ctx.buffer );
                             ctx.remainingChunkSize = partialChunkSize << 8;
                             return IN_HEADER.readChunkSize( ctx );
                         }
@@ -220,7 +229,7 @@ public State readChunkSize( BufferingChunkedInput ctx ) throws IOException
                         {
                             //Now we have enough space to read the rest of the chunk size
                             byte partialChunkSize = ctx.buffer.get();
-                            ctx.remainingChunkSize = (ctx.remainingChunkSize | partialChunkSize) & 0xFFFF;
+                            ctx.remainingChunkSize = ctx.remainingChunkSize | (partialChunkSize & 0xFF);
                             return IN_CHUNK;
                         }
                         else
@@ -383,6 +392,12 @@ public byte peekByte() throws IOException
         return buffer.get( buffer.position() );
     }
 
+    static int getUnsignedByteFromBuffer( ByteBuffer buffer )
+    {
+        return buffer.get() & 0xFF;
+    }
+
+
     private boolean hasMoreDataUnreadInCurrentChunk()
     {
         return remainingChunkSize > 0;
diff --git a/driver/src/test/java/org/neo4j/driver/internal/connector/socket/BufferingChunkedInputTest.java b/driver/src/test/java/org/neo4j/driver/internal/connector/socket/BufferingChunkedInputTest.java
@@ -92,6 +92,47 @@ public void shouldReadOneByteWhenSplitHeader() throws IOException
         assertThat( b2, equalTo( (byte) 37 ) );
     }
 
+    @Test
+    public void shouldReadChunkWithSplitHeaderForBigMessages() throws IOException
+    {
+        // Given
+        int packetSize = 384;
+        BufferingChunkedInput input =
+                new BufferingChunkedInput( packets( packet( 1 ), packet( -128 ), fillPacket( packetSize, 1 ) ) );
+
+        // Then
+        assertThat( input.readByte(), equalTo( (byte) 1 ) );
+        assertThat( input.remainingChunkSize(), equalTo( packetSize - 1 ) );
+
+        for ( int i = 1; i < packetSize; i++ )
+        {
+            assertThat( input.readByte(), equalTo( (byte) 1 ) );
+        }
+        assertThat( input.remainingChunkSize(), equalTo( 0 ) );
+    }
+
+    @Test
+    public void shouldReadChunkWithSplitHeaderForBigMessagesWhenInternalBufferHasOneByte() throws IOException
+    {
+        // Given
+        int packetSize = 32780;
+        BufferingChunkedInput input =
+                new BufferingChunkedInput( packets( packet( -128 ), packet( 12 ), fillPacket( packetSize, 1 ) ), 1);
+
+        // Then
+        assertThat( input.readByte(), equalTo( (byte) 1 ) );
+        assertThat( input.remainingChunkSize(), equalTo( packetSize - 1 ) );
+    }
+
+    @Test
+    public void shouldReadUnsignedByteFromBuffer() throws IOException
+    {
+        ByteBuffer buffer = ByteBuffer.allocate( 1 );
+        buffer.put( (byte) -1 );
+        buffer.flip();
+        assertThat(BufferingChunkedInput.getUnsignedByteFromBuffer( buffer ), equalTo( 255 ));
+    }
+
     @Test
     public void shouldReadOneByteInOneChunkWhenBustingBuffer() throws IOException
     {
@@ -415,7 +456,7 @@ public void close() throws IOException
         BufferingChunkedInput input = new BufferingChunkedInput( channel );
 
         // Then
-        assertThat(input.readByte(), equalTo( (byte)11 ));
+        assertThat( input.readByte(), equalTo( (byte) 11 ) );
 
     }
 
@@ -435,6 +476,17 @@ public void shouldFailNicelyOnClosedConnections() throws IOException
         input.readByte();
     }
 
+    private ReadableByteChannel fillPacket( int size, int value )
+    {
+        int[] ints = new int[size];
+        for ( int i = 0; i < size; i++ )
+        {
+            ints[i] = value;
+        }
+
+        return packet( ints );
+    }
+
     private ReadableByteChannel packet( int... bytes )
     {
         byte[] byteArray = new byte[bytes.length];
diff --git a/driver/src/test/java/org/neo4j/driver/internal/connector/socket/ChunkedInputTest.java b/driver/src/test/java/org/neo4j/driver/internal/connector/socket/ChunkedInputTest.java
@@ -18,6 +18,11 @@
  */
 package org.neo4j.driver.internal.connector.socket;
 
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.mockito.Matchers;
+
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -26,11 +31,6 @@
 import java.nio.channels.ReadableByteChannel;
 import java.util.Arrays;
 
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-import org.mockito.Matchers;
-
 import org.neo4j.driver.v1.exceptions.ClientException;
 import org.neo4j.driver.v1.util.RecordingByteChannel;
 
