diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 1aec0d46..2d42ae2f 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -2,7 +2,6 @@ package buf import ( "io" - "net" "v2ray.com/core/common" "v2ray.com/core/common/errors" @@ -204,15 +203,6 @@ func (mb MultiBuffer) String() string { return serial.Concat(v...) } -// ToNetBuffers converts this MultiBuffer to net.Buffers. The return net.Buffers points to the same content of the MultiBuffer. -func (mb MultiBuffer) ToNetBuffers() net.Buffers { - bs := make([][]byte, len(mb)) - for i, b := range mb { - bs[i] = b.Bytes() - } - return bs -} - // SliceBySize splits the beginning of this MultiBuffer into another one, for at most size bytes. func (mb *MultiBuffer) SliceBySize(size int32) MultiBuffer { slice := NewMultiBufferCap(10) diff --git a/common/buf/writer.go b/common/buf/writer.go index 49c112ef..51b10464 100644 --- a/common/buf/writer.go +++ b/common/buf/writer.go @@ -2,6 +2,7 @@ package buf import ( "io" + "net" "v2ray.com/core/common" "v2ray.com/core/common/errors" @@ -10,6 +11,8 @@ import ( // BufferToBytesWriter is a Writer that writes alloc.Buffer into underlying writer. type BufferToBytesWriter struct { io.Writer + + cache [][]byte } // WriteMultiBuffer implements Writer. This method takes ownership of the given buffer. @@ -25,10 +28,16 @@ func (w *BufferToBytesWriter) WriteMultiBuffer(mb MultiBuffer) error { return WriteAllBytes(w.Writer, mb[0].Bytes()) } - bs := mb.ToNetBuffers() + bs := w.cache + for _, b := range mb { + bs = append(bs, b.Bytes()) + } + w.cache = bs[:0] + + nb := net.Buffers(bs) for size > 0 { - n, err := bs.WriteTo(w.Writer) + n, err := nb.WriteTo(w.Writer) if err != nil { return err } diff --git a/common/crypto/io.go b/common/crypto/io.go index c8c044f0..19a3a25c 100644 --- a/common/crypto/io.go +++ b/common/crypto/io.go @@ -32,44 +32,35 @@ var ( ) type CryptionWriter struct { - stream cipher.Stream - writer io.Writer + stream cipher.Stream + writer io.Writer + bufWriter buf.Writer } // NewCryptionWriter creates a new CryptionWriter. func NewCryptionWriter(stream cipher.Stream, writer io.Writer) *CryptionWriter { return &CryptionWriter{ - stream: stream, - writer: writer, + stream: stream, + writer: writer, + bufWriter: buf.NewWriter(writer), } } // Write implements io.Writer.Write(). func (w *CryptionWriter) Write(data []byte) (int, error) { w.stream.XORKeyStream(data, data) - return w.writer.Write(data) + + if err := buf.WriteAllBytes(w.writer, data); err != nil { + return 0, err + } + return len(data), nil } // WriteMultiBuffer implements buf.Writer. func (w *CryptionWriter) WriteMultiBuffer(mb buf.MultiBuffer) error { - defer mb.Release() - - size := mb.Len() - if size == 0 { - return nil - } - - bs := mb.ToNetBuffers() - for _, b := range bs { - w.stream.XORKeyStream(b, b) + for _, b := range mb { + w.stream.XORKeyStream(b.Bytes(), b.Bytes()) } - for size > 0 { - n, err := bs.WriteTo(w.writer) - if err != nil { - return err - } - size -= int32(n) - } - return nil + return w.bufWriter.WriteMultiBuffer(mb) }