From 17c368f9c243eb18940dfaaa90bb04b93f52054b Mon Sep 17 00:00:00 2001
From: OlivierSalasc <olivier.salasc@gmail.com>
Date: Thu, 27 Sep 2018 14:49:25 +0200
Subject: [PATCH] add Buffer capacity check for type UInt64 (#454)

---
 .../java/net/schmizz/sshj/common/Buffer.java  |  1 +
 .../net/schmizz/sshj/common/BufferTest.java   | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/main/java/net/schmizz/sshj/common/Buffer.java b/src/main/java/net/schmizz/sshj/common/Buffer.java
index 183de8d3..c11852d7 100644
--- 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 class Buffer<T extends Buffer<T>> {
 
     @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
index 98a7a203..7e61ad9a 100644
--- 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 class BufferTest {
             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);
+    }
+
 }