From e5525715fb3726646daa610e7f55a7128efa00b5 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 16 Apr 2017 22:30:29 +0200
Subject: [PATCH] optimize cfb stream writer

---
 common/buf/multi_buffer.go | 13 ++++++++++++-
 common/buf/writer.go       | 14 ++++----------
 common/crypto/io.go        | 23 +++++++++++++++++------
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go
index e6809cf38..4bc7a93e6 100644
--- a/common/buf/multi_buffer.go
+++ b/common/buf/multi_buffer.go
@@ -1,6 +1,9 @@
 package buf
 
-import "io"
+import (
+	"io"
+	"net"
+)
 
 type MultiBufferWriter interface {
 	WriteMultiBuffer(MultiBuffer) (int, error)
@@ -75,3 +78,11 @@ func (mb MultiBuffer) Release() {
 		mb[i] = nil
 	}
 }
+
+func (mb MultiBuffer) ToNetBuffers() net.Buffers {
+	bs := make([][]byte, len(mb))
+	for i, b := range mb {
+		bs[i] = b.Bytes()
+	}
+	return bs
+}
diff --git a/common/buf/writer.go b/common/buf/writer.go
index 141c8b354..46d528e68 100644
--- a/common/buf/writer.go
+++ b/common/buf/writer.go
@@ -1,9 +1,6 @@
 package buf
 
-import (
-	"io"
-	"net"
-)
+import "io"
 
 // BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer.
 type BufferToBytesWriter struct {
@@ -18,12 +15,9 @@ func (v *BufferToBytesWriter) Write(mb MultiBuffer) error {
 	}
 
 	defer mb.Release()
-	bs := make([][]byte, len(mb))
-	for i, b := range mb {
-		bs[i] = b.Bytes()
-	}
-	nbs := net.Buffers(bs)
-	_, err := nbs.WriteTo(v.writer)
+
+	bs := mb.ToNetBuffers()
+	_, err := bs.WriteTo(v.writer)
 	return err
 }
 
diff --git a/common/crypto/io.go b/common/crypto/io.go
index fe9d29726..133094f7b 100644
--- a/common/crypto/io.go
+++ b/common/crypto/io.go
@@ -3,6 +3,8 @@ package crypto
 import (
 	"crypto/cipher"
 	"io"
+
+	"v2ray.com/core/common/buf"
 )
 
 type CryptionReader struct {
@@ -17,10 +19,10 @@ func NewCryptionReader(stream cipher.Stream, reader io.Reader) *CryptionReader {
 	}
 }
 
-func (v *CryptionReader) Read(data []byte) (int, error) {
-	nBytes, err := v.reader.Read(data)
+func (r *CryptionReader) Read(data []byte) (int, error) {
+	nBytes, err := r.reader.Read(data)
 	if nBytes > 0 {
-		v.stream.XORKeyStream(data[:nBytes], data[:nBytes])
+		r.stream.XORKeyStream(data[:nBytes], data[:nBytes])
 	}
 	return nBytes, err
 }
@@ -39,7 +41,16 @@ func NewCryptionWriter(stream cipher.Stream, writer io.Writer) *CryptionWriter {
 }
 
 // Write implements io.Writer.Write().
-func (v *CryptionWriter) Write(data []byte) (int, error) {
-	v.stream.XORKeyStream(data, data)
-	return v.writer.Write(data)
+func (w *CryptionWriter) Write(data []byte) (int, error) {
+	w.stream.XORKeyStream(data, data)
+	return w.writer.Write(data)
+}
+
+func (w *CryptionWriter) WriteMultiBuffer(mb buf.MultiBuffer) (int, error) {
+	bs := mb.ToNetBuffers()
+	for _, b := range bs {
+		w.stream.XORKeyStream(b, b)
+	}
+	nBytes, err := bs.WriteTo(w.writer)
+	return int(nBytes), err
 }