optimize cfb stream writer

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
+}

common/buf/writer.go

@@ -1,9 +1,6 @@
 package buf
 
-import (
+import "io"
-	"io"
-	"net"
-)
 
 // 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 := mb.ToNetBuffers()
-		bs[i] = b.Bytes()
+	_, err := bs.WriteTo(v.writer)
-	}
-	nbs := net.Buffers(bs)
-	_, err := nbs.WriteTo(v.writer)
 	return err
 }
 

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) {
+func (r *CryptionReader) Read(data []byte) (int, error) {
-	nBytes, err := v.reader.Read(data)
+	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) {
+func (w *CryptionWriter) Write(data []byte) (int, error) {
-	v.stream.XORKeyStream(data, data)
+	w.stream.XORKeyStream(data, data)
-	return v.writer.Write(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
 }