diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index e6809cf3..4bc7a93e 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 141c8b35..46d528e6 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 fe9d2972..133094f7 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 }