From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 12 Sep 2017 15:00:43 +0000 (+1000)
Subject: mse: Reuse writer write buffer
X-Git-Tag: v1.0.0~403
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=7e056f7753c5d78d131bdef7f1ba088d4bee0c53;p=btrtrc.git

mse: Reuse writer write buffer
---

diff --git a/mse/mse.go b/mse/mse.go
index 02f649a6..421ade3f 100644
--- a/mse/mse.go
+++ b/mse/mse.go
@@ -116,17 +116,29 @@ func newCipherReader(c *rc4.Cipher, r io.Reader) io.Reader {
 type cipherWriter struct {
 	c *rc4.Cipher
 	w io.Writer
+	b []byte
 }
 
 func (cr *cipherWriter) Write(b []byte) (n int, err error) {
-	be := make([]byte, len(b))
-	cr.c.XORKeyStream(be, b)
-	n, err = cr.w.Write(be)
-	if n != len(be) {
+	be := func() []byte {
+		if len(cr.b) < len(b) {
+			return make([]byte, len(b))
+		} else {
+			ret := cr.b
+			cr.b = nil
+			return ret
+		}
+	}()
+	cr.c.XORKeyStream(be[:], b)
+	n, err = cr.w.Write(be[:len(b)])
+	if n != len(b) {
 		// The cipher will have advanced beyond the callers stream position.
 		// We can't use the cipher anymore.
 		cr.c = nil
 	}
+	if len(be) > len(cr.b) {
+		cr.b = be
+	}
 	return
 }
 
@@ -390,7 +402,7 @@ func (h *handshake) initerSteps() (ret io.ReadWriter, err error) {
 	if err != nil {
 		return
 	}
-	ret = readWriter{r, &cipherWriter{e, h.conn}}
+	ret = readWriter{r, &cipherWriter{e, h.conn, nil}}
 	return
 }
 
@@ -449,7 +461,7 @@ func (h *handshake) receiverSteps() (ret io.ReadWriter, err error) {
 		unmarshal(r, h.ia)
 	}
 	buf := &bytes.Buffer{}
-	w := cipherWriter{h.newEncrypt(false), buf}
+	w := cipherWriter{h.newEncrypt(false), buf, nil}
 	padLen = uint16(newPadLen())
 	err = marshal(&w, &vc, uint32(cryptoMethodRC4), padLen, zeroPad[:padLen])
 	if err != nil {
@@ -459,7 +471,7 @@ func (h *handshake) receiverSteps() (ret io.ReadWriter, err error) {
 	if err != nil {
 		return
 	}
-	ret = readWriter{io.MultiReader(bytes.NewReader(h.ia), r), &cipherWriter{w.c, h.conn}}
+	ret = readWriter{io.MultiReader(bytes.NewReader(h.ia), r), &cipherWriter{w.c, h.conn, nil}}
 	return
 }