add Buffer capacity check for type UInt64 (#454)

hierynomus · Sep 27, 2018 · 17c368f · 17c368f
1 parent 4de9f8a
commit 17c368f
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/src/main/java/net/schmizz/sshj/common/Buffer.java b/src/main/java/net/schmizz/sshj/common/Buffer.java
@@ -372,6 +372,7 @@ public T putUInt64(BigInteger uint64) {
 
     @SuppressWarnings("unchecked")
     private T putUInt64Unchecked(long uint64) {
+        ensureCapacity(8);
         data[wpos++] = (byte) (uint64 >> 56);
         data[wpos++] = (byte) (uint64 >> 48);
         data[wpos++] = (byte) (uint64 >> 40);

diff --git a/src/test/java/net/schmizz/sshj/common/BufferTest.java b/src/test/java/net/schmizz/sshj/common/BufferTest.java
@@ -146,4 +146,28 @@ public void shouldHaveSameUInt64EncodingForBigIntegerAndLong() {
             assertArrayEquals("Value: " + value, bytesLong, bytesBigInt);
         }
     }
+
+
+    @Test
+    public void shouldExpandCapacityOfUInt32(){
+        PlainBuffer buf = new PlainBuffer();
+        for(int i=0;i<Buffer.DEFAULT_SIZE+1;i+=4) {
+            buf.putUInt32(1l);
+        }
+        /* Buffer should have been expanded at this point*/
+        assertEquals(Buffer.DEFAULT_SIZE*2,buf.data.length);
+    }
+
+    @Test
+    public void shouldExpandCapacityOfUInt64(){
+        BigInteger bigUint64 = BigInteger.valueOf(Long.MAX_VALUE);
+        PlainBuffer buf = new PlainBuffer();
+        assertEquals(Buffer.DEFAULT_SIZE,buf.data.length);
+        for(int i=0;i<Buffer.DEFAULT_SIZE+1;i+=8) {
+            buf.putUInt64(bigUint64.longValue());
+        }
+        /* Buffer should have been expanded at this point*/
+        assertEquals(Buffer.DEFAULT_SIZE*2,buf.data.length);
+    }
+
 }