From e6d857999c03294c80b641e8f70cfa6fa4db5a47 Mon Sep 17 00:00:00 2001
From: Riobard <me@riobard.com>
Date: Sat, 18 Feb 2017 23:41:55 +0800
Subject: [PATCH] Reduce allocations

---
 shadowaead/packet.go   | 21 ++++++++++++---------
 shadowstream/packet.go | 10 +++++++---
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/shadowaead/packet.go b/shadowaead/packet.go
index 18a6f8d1..d8f20a40 100644
--- a/shadowaead/packet.go
+++ b/shadowaead/packet.go
@@ -5,11 +5,14 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 )
 
 // ErrShortPacket means that the packet is too short for a valid encrypted packet.
 var ErrShortPacket = errors.New("short packet")
 
+var _zerononce [128]byte // read-only. 128 bytes is more than enough.
+
 // Pack encrypts plaintext using Cipher with a randomly generated salt and
 // returns a slice of dst containing the encrypted packet and any error occurred.
 // Ensure len(dst) >= ciph.SaltSize() + len(plaintext) + aead.Overhead().
@@ -28,8 +31,7 @@ func Pack(dst, plaintext []byte, ciph Cipher) ([]byte, error) {
 	if len(dst) < saltSize+len(plaintext)+aead.Overhead() {
 		return nil, io.ErrShortBuffer
 	}
-	nonce := make([]byte, aead.NonceSize())
-	b := aead.Seal(dst[saltSize:saltSize], nonce, plaintext, nil)
+	b := aead.Seal(dst[saltSize:saltSize], _zerononce[:aead.NonceSize()], plaintext, nil)
 	return dst[:saltSize+len(b)], nil
 }
 
@@ -51,27 +53,28 @@ func Unpack(dst, pkt []byte, ciph Cipher) ([]byte, error) {
 	if saltSize+len(dst)+aead.Overhead() < len(pkt) {
 		return nil, io.ErrShortBuffer
 	}
-
-	nonce := make([]byte, aead.NonceSize())
-	b, err := aead.Open(dst[:0], nonce, pkt[saltSize:], nil)
+	b, err := aead.Open(dst[:0], _zerononce[:aead.NonceSize()], pkt[saltSize:], nil)
 	return b, err
 }
 
 type packetConn struct {
 	net.PacketConn
 	Cipher
+	sync.Mutex
+	buf []byte // write lock
 }
 
 // NewPacketConn wraps a net.PacketConn with cipher
 func NewPacketConn(c net.PacketConn, ciph Cipher) net.PacketConn {
-	return &packetConn{PacketConn: c, Cipher: ciph}
+	const maxPacketSize = 64 * 1024
+	return &packetConn{PacketConn: c, Cipher: ciph, buf: make([]byte, maxPacketSize)}
 }
 
 // WriteTo encrypts b and write to addr using the embedded PacketConn.
 func (c *packetConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	const overhead = 16
-	buf := make([]byte, c.Cipher.SaltSize()+len(b)+overhead)
-	buf, err := Pack(buf, b, c)
+	c.Lock()
+	defer c.Unlock()
+	buf, err := Pack(c.buf, b, c)
 	if err != nil {
 		return 0, err
 	}
diff --git a/shadowstream/packet.go b/shadowstream/packet.go
index 31d2bc80..8bbb27b5 100644
--- a/shadowstream/packet.go
+++ b/shadowstream/packet.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 )
 
 // ErrShortPacket means the packet is too short to be a valid encrypted packet.
@@ -45,16 +46,19 @@ func Unpack(dst, pkt []byte, s Cipher) ([]byte, error) {
 type packetConn struct {
 	net.PacketConn
 	Cipher
+	buf        []byte
+	sync.Mutex // write lock
 }
 
 // NewPacketConn wraps a net.PacketConn with stream cipher encryption/decryption.
 func NewPacketConn(c net.PacketConn, ciph Cipher) net.PacketConn {
-	return &packetConn{PacketConn: c, Cipher: ciph}
+	return &packetConn{PacketConn: c, Cipher: ciph, buf: make([]byte, 64*1024)}
 }
 
 func (c *packetConn) WriteTo(b []byte, addr net.Addr) (int, error) {
-	buf := make([]byte, c.IVSize()+len(b))
-	_, err := Pack(buf, b, c.Cipher)
+	c.Lock()
+	defer c.Unlock()
+	buf, err := Pack(c.buf, b, c.Cipher)
 	if err != nil {
 		return 0, err
 	}