From 5c4e33f7593cb1ac491c7a3a352890320496767a Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 18 Nov 2018 20:16:14 +0100
Subject: [PATCH] rewrite SliceBySize

---
 common/buf/multi_buffer.go      | 57 +++++++++++++++++++--------------
 common/buf/multi_buffer_test.go |  2 +-
 common/buf/reader.go            |  3 +-
 common/crypto/chunk.go          |  3 +-
 common/mux/writer.go            |  2 +-
 5 files changed, 39 insertions(+), 28 deletions(-)

diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go
index 43ea2b85f..91b9a2479 100644
--- a/common/buf/multi_buffer.go
+++ b/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
diff --git a/common/buf/multi_buffer_test.go b/common/buf/multi_buffer_test.go
index c33bd856c..1e4125225 100644
--- a/common/buf/multi_buffer_test.go
+++ b/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))
 }
diff --git a/common/buf/reader.go b/common/buf/reader.go
index d07111c7b..1d21b8a87 100644
--- a/common/buf/reader.go
+++ b/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
 	}
diff --git a/common/crypto/chunk.go b/common/crypto/chunk.go
index 65c5c0f82..164336d0f 100755
--- a/common/crypto/chunk.go
+++ b/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()))
diff --git a/common/mux/writer.go b/common/mux/writer.go
index 5259bedbe..f3b3c1d35 100644
--- a/common/mux/writer.go
+++ b/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