From 86a6a1998f1d2bf63262718644c1ef3a77fa8d28 Mon Sep 17 00:00:00 2001
From: Kale Blankenship <kale@lemnisys.com>
Date: Tue, 26 Nov 2019 16:18:31 -0800
Subject: [PATCH] reclaim receive buffer space when more than maxFrameSize is
 free

* Prevents unbounded buffer growth when more than one frame is read
  into the buffer.
* Change buffer to grow by 2x rather than 512 bytes.

Updates #191
---
 buffer.go | 15 ++++++++++++++-
 conn.go   |  8 +++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/buffer.go b/buffer.go
index 1cbd095f..6dd81937 100644
--- a/buffer.go
+++ b/buffer.go
@@ -32,6 +32,15 @@ func (b *buffer) reset() {
 	b.i = 0
 }
 
+// reclaim shifts used buffer space to the beginning of the
+// underlying slice.
+func (b *buffer) reclaim() {
+	l := b.len()
+	copy(b.b[:l], b.b[b.i:])
+	b.b = b.b[:l]
+	b.i = 0
+}
+
 func (b *buffer) readCheck(n int64) bool {
 	return int64(b.i)+n > int64(len(b.b))
 }
@@ -94,7 +103,11 @@ func (b *buffer) readFromOnce(r io.Reader) error {
 
 	l := len(b.b)
 	if cap(b.b)-l < minRead {
-		new := make([]byte, l, l+minRead)
+		total := l * 2
+		if total == 0 {
+			total = minRead
+		}
+		new := make([]byte, l, total)
 		copy(new, b.b)
 		b.b = new
 	}
diff --git a/conn.go b/conn.go
index b46864b5..9f6c5531 100644
--- a/conn.go
+++ b/conn.go
@@ -446,8 +446,14 @@ func (c *conn) connReader() {
 	)
 
 	for {
-		if buf.len() == 0 {
+		switch {
+		// Cheaply reuse free buffer space when fully read.
+		case buf.len() == 0:
 			buf.reset()
+
+		// Prevent excessive/unbounded growth by shifting data to beginning of buffer.
+		case int64(buf.i) > int64(c.maxFrameSize):
+			buf.reclaim()
 		}
 
 		// need to read more if buf doesn't contain the complete frame