rewrite SliceBySize

common/buf/multi_buffer.go

@@ -3,7 +3,6 @@ package buf
 import (
 	"io"
 
-	"v2ray.com/core/common"
 	"v2ray.com/core/common/errors"
 	"v2ray.com/core/common/serial"
 )
@@ -127,6 +126,39 @@ func SplitFirst(mb MultiBuffer) (MultiBuffer, *Buffer) {
 	return mb, b
 }
 
+// SplitSize splits the beginning of the MultiBuffer into another one, for at most size bytes.
+func SplitSize(mb MultiBuffer, size int32) (MultiBuffer, MultiBuffer) {
+	if len(mb) == 0 {
+		return mb, nil
+	}
+
+	if mb[0].Len() > size {
+		b := New()
+		copy(b.Extend(size), mb[0].BytesTo(size))
+		mb[0].Advance(size)
+		return mb, MultiBuffer{b}
+	}
+
+	totalBytes := int32(0)
+	var r MultiBuffer
+	endIndex := 0
+	for i := range mb {
+		if totalBytes+mb[i].Len() > size {
+			endIndex = i
+			break
+		}
+		r = append(r, mb[i])
+		mb[i] = nil
+	}
+	if endIndex == len(mb) {
+		// To reuse mb array
+		mb = mb[:0]
+	} else {
+		mb = mb[endIndex:]
+	}
+	return mb, r
+}
+
 // Len returns the total number of bytes in the MultiBuffer.
 func (mb MultiBuffer) Len() int32 {
 	if mb == nil {
@@ -159,29 +191,6 @@ func (mb MultiBuffer) String() string {
 	return serial.Concat(v...)
 }
 
-// SliceBySize splits the beginning of this MultiBuffer into another one, for at most size bytes.
-func (mb *MultiBuffer) SliceBySize(size int32) MultiBuffer {
-	slice := make(MultiBuffer, 0, 10)
-	sliceSize := int32(0)
-	endIndex := len(*mb)
-	for i, b := range *mb {
-		if b.Len()+sliceSize > size {
-			endIndex = i
-			break
-		}
-		sliceSize += b.Len()
-		slice = append(slice, b)
-		(*mb)[i] = nil
-	}
-	*mb = (*mb)[endIndex:]
-	if endIndex == 0 && len(*mb) > 0 {
-		b := New()
-		common.Must2(b.ReadFullFrom((*mb)[0], size))
-		return MultiBuffer{b}
-	}
-	return slice
-}
-
 // MultiBufferContainer is a ReadWriteCloser wrapper over MultiBuffer.
 type MultiBufferContainer struct {
 	MultiBuffer

common/buf/multi_buffer_test.go

@@ -44,6 +44,6 @@ func TestMultiBufferSliceBySizeLarge(t *testing.T) {
 
 	mb := MergeBytes(nil, lb)
 
-	mb2 := mb.SliceBySize(1024)
+	mb, mb2 := SplitSize(mb, 1024)
 	assert(mb2.Len(), Equals, int32(1024))
 }

common/buf/reader.go

@@ -87,7 +87,8 @@ func (r *BufferedReader) ReadAtMost(size int32) (MultiBuffer, error) {
 		r.Buffer = mb
 	}
 
-	mb := r.Buffer.SliceBySize(size)
+	rb, mb := SplitSize(r.Buffer, size)
+	r.Buffer = rb
 	if r.Buffer.IsEmpty() {
 		r.Buffer = nil
 	}

common/crypto/chunk.go

@@ -143,7 +143,8 @@ func (w *ChunkStreamWriter) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	mb2Write := make(buf.MultiBuffer, 0, mbLen/buf.Size+mbLen/sliceSize+2)
 
 	for {
-		slice := mb.SliceBySize(sliceSize)
+		mb2, slice := buf.SplitSize(mb, sliceSize)
+		mb = mb2
 
 		b := buf.New()
 		w.sizeEncoder.Encode(uint16(slice.Len()), b.Extend(w.sizeEncoder.SizeBytes()))

common/mux/writer.go

@@ -94,7 +94,7 @@ func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
 	for !mb.IsEmpty() {
 		var chunk buf.MultiBuffer
 		if w.transferType == protocol.TransferTypeStream {
-			chunk = mb.SliceBySize(8 * 1024)
+			mb, chunk = buf.SplitSize(mb, 8*1024)
 		} else {
 			mb2, b := buf.SplitFirst(mb)
 			mb = mb2