Browse Source

remove ToNetBuffers

Darien Raymond 7 năm trước cách đây
mục cha
commit
2594f7027a
3 tập tin đã thay đổi với 25 bổ sung35 xóa
  1. 0 10
      common/buf/multi_buffer.go
  2. 11 2
      common/buf/writer.go
  3. 14 23
      common/crypto/io.go

+ 0 - 10
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)

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

+ 14 - 23
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)
 }