From bcd5d026fe8be028722d537416038eabba017e78 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Sun, 18 Nov 2018 19:57:29 +0100
Subject: [PATCH] rewrite SplitFirst

---
 common/buf/multi_buffer.go | 32 +++++++++++++++++++++-----------
 common/crypto/auth_test.go |  4 ++--
 common/mux/writer.go       |  4 +++-
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go
index 9cd31dc3..43ea2b85 100644
--- a/common/buf/multi_buffer.go
+++ b/common/buf/multi_buffer.go
@@ -35,6 +35,7 @@ func MergeMulti(dest MultiBuffer, src MultiBuffer) (MultiBuffer, MultiBuffer) {
 	return dest, src[:0]
 }
 
+// MergeBytes merges the given bytes into MultiBuffer and return the new address of the merged MultiBuffer.
 func MergeBytes(dest MultiBuffer, src []byte) MultiBuffer {
 	n := len(dest)
 	if n > 0 && !(dest)[n-1].IsFull() {
@@ -94,6 +95,8 @@ func ReadFrom(reader io.Reader) (MultiBuffer, error) {
 	}
 }
 
+// SplitBytes splits the given amount of bytes from the beginning of the MultiBuffer.
+// It returns the new address of MultiBuffer leftover, and number of bytes written into the input byte slice.
 func SplitBytes(mb MultiBuffer, b []byte) (MultiBuffer, int) {
 	totalBytes := 0
 
@@ -112,6 +115,18 @@ func SplitBytes(mb MultiBuffer, b []byte) (MultiBuffer, int) {
 	return mb, totalBytes
 }
 
+// SplitFirst splits the first Buffer from the beginning of the MultiBuffer.
+func SplitFirst(mb MultiBuffer) (MultiBuffer, *Buffer) {
+	if len(mb) == 0 {
+		return mb, nil
+	}
+
+	b := mb[0]
+	mb[0] = nil
+	mb = mb[1:]
+	return mb, b
+}
+
 // Len returns the total number of bytes in the MultiBuffer.
 func (mb MultiBuffer) Len() int32 {
 	if mb == nil {
@@ -167,21 +182,12 @@ func (mb *MultiBuffer) SliceBySize(size int32) MultiBuffer {
 	return slice
 }
 
-// SplitFirst splits out the first Buffer in this MultiBuffer.
-func (mb *MultiBuffer) SplitFirst() *Buffer {
-	if len(*mb) == 0 {
-		return nil
-	}
-	b := (*mb)[0]
-	(*mb)[0] = nil
-	*mb = (*mb)[1:]
-	return b
-}
-
+// MultiBufferContainer is a ReadWriteCloser wrapper over MultiBuffer.
 type MultiBufferContainer struct {
 	MultiBuffer
 }
 
+// Read implements io.Reader.
 func (c *MultiBufferContainer) Read(b []byte) (int, error) {
 	if c.MultiBuffer.IsEmpty() {
 		return 0, io.EOF
@@ -192,23 +198,27 @@ func (c *MultiBufferContainer) Read(b []byte) (int, error) {
 	return nBytes, nil
 }
 
+// ReadMultiBuffer implements Reader.
 func (c *MultiBufferContainer) ReadMultiBuffer() (MultiBuffer, error) {
 	mb := c.MultiBuffer
 	c.MultiBuffer = nil
 	return mb, nil
 }
 
+// Write implements io.Writer.
 func (c *MultiBufferContainer) Write(b []byte) (int, error) {
 	c.MultiBuffer = MergeBytes(c.MultiBuffer, b)
 	return len(b), nil
 }
 
+// WriteMultiBuffer implement Writer.
 func (c *MultiBufferContainer) WriteMultiBuffer(b MultiBuffer) error {
 	mb, _ := MergeMulti(c.MultiBuffer, b)
 	c.MultiBuffer = mb
 	return nil
 }
 
+// Close implement io.Closer.
 func (c *MultiBufferContainer) Close() error {
 	c.MultiBuffer = ReleaseMulti(c.MultiBuffer)
 	return nil
diff --git a/common/crypto/auth_test.go b/common/crypto/auth_test.go
index b79a15be..d06f5b01 100644
--- a/common/crypto/auth_test.go
+++ b/common/crypto/auth_test.go
@@ -116,10 +116,10 @@ func TestAuthenticationReaderWriterPacket(t *testing.T) {
 	mb, err := reader.ReadMultiBuffer()
 	assert(err, IsNil)
 
-	b1 := mb.SplitFirst()
+	mb, b1 := buf.SplitFirst(mb)
 	assert(b1.String(), Equals, "abcd")
 
-	b2 := mb.SplitFirst()
+	mb, b2 := buf.SplitFirst(mb)
 	assert(b2.String(), Equals, "efgh")
 
 	assert(mb.IsEmpty(), IsTrue)
diff --git a/common/mux/writer.go b/common/mux/writer.go
index 2dc3904c..5259bedb 100644
--- a/common/mux/writer.go
+++ b/common/mux/writer.go
@@ -96,7 +96,9 @@ func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
 		if w.transferType == protocol.TransferTypeStream {
 			chunk = mb.SliceBySize(8 * 1024)
 		} else {
-			chunk = buf.MultiBuffer{mb.SplitFirst()}
+			mb2, b := buf.SplitFirst(mb)
+			mb = mb2
+			chunk = buf.MultiBuffer{b}
 		}
 		if err := w.writeData(chunk); err != nil {
 			return err