From 0c16ab6b32f2fa78b137985915b08953d589296a Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Sat, 24 Feb 2018 19:29:35 -0800
Subject: [PATCH] brontide: reduce memory allocs by using static buf for next
 header+msg
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

In this commit, we reduce the total number of allocations that a
brontide session will incur over its lifetime. Profiling on one of my
nodes showed that we were generating a lot of garbage due to
re-creating a 65KB buffer to read the next message each time the
ReadMessage method was called.

To reduce the total number of memory allocations, we’ll now simply
re-use a buffer for both the cipher text header, and the cipher text
itself.
---
 brontide/noise.go | 26 ++++++++++++++++++++------
 1 file changed, 20 insertions(+), 6 deletions(-)

diff --git a/brontide/noise.go b/brontide/noise.go
index 723b7bfb..2bc30969 100644
--- a/brontide/noise.go
+++ b/brontide/noise.go
@@ -345,6 +345,20 @@ type Machine struct {
 	ephemeralGen func() (*btcec.PrivateKey, error)
 
 	handshakeState
+
+	// nextCipherHeader is a static buffer that we'll use to read in the
+	// next ciphertext header from the wire. The header is a 2 byte length
+	// (of the next ciphertext), followed by a 16 byte MAC.
+	nextCipherHeader [lengthHeaderSize + macSize]byte
+
+	// nextCipherText is a static buffer that we'll use to read in the
+	// bytes of the next cipher text message. As all messages in the
+	// protocol MUST be below 65KB plus our macSize, this will be
+	// sufficient to buffer all messages from the socket when we need to
+	// read the next one. Having a fixed buffer that's re-used also means
+	// that we save on allocations as we don't need to create a new one
+	// each time.
+	nextCipherText [math.MaxUint16 + macSize]byte
 }
 
 // NewBrontideMachine creates a new instance of the brontide state-machine. If
@@ -698,25 +712,25 @@ func (b *Machine) WriteMessage(w io.Writer, p []byte) error {
 // ReadMessage attempts to read the next message from the passed io.Reader. In
 // the case of an authentication error, a non-nil error is returned.
 func (b *Machine) ReadMessage(r io.Reader) ([]byte, error) {
-	var cipherLen [lengthHeaderSize + macSize]byte
-	if _, err := io.ReadFull(r, cipherLen[:]); err != nil {
+	if _, err := io.ReadFull(r, b.nextCipherHeader[:]); err != nil {
 		return nil, err
 	}
 
 	// Attempt to decrypt+auth the packet length present in the stream.
-	pktLenBytes, err := b.recvCipher.Decrypt(nil, nil, cipherLen[:])
+	pktLenBytes, err := b.recvCipher.Decrypt(
+		nil, nil, b.nextCipherHeader[:],
+	)
 	if err != nil {
 		return nil, err
 	}
 
 	// Next, using the length read from the packet header, read the
 	// encrypted packet itself.
-	var cipherText [math.MaxUint16 + macSize]byte
 	pktLen := uint32(binary.BigEndian.Uint16(pktLenBytes)) + macSize
-	if _, err := io.ReadFull(r, cipherText[:pktLen]); err != nil {
+	if _, err := io.ReadFull(r, b.nextCipherText[:pktLen]); err != nil {
 		return nil, err
 	}
 
 	// TODO(roasbeef): modify to let pass in slice
-	return b.recvCipher.Decrypt(nil, nil, cipherText[:pktLen])
+	return b.recvCipher.Decrypt(nil, nil, b.nextCipherText[:pktLen])
 }