diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 43ea2b85..91b9a247 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 c33bd856..1e412522 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 d07111c7..1d21b8a8 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 65c5c0f8..164336d0 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 5259bedb..f3b3c1d3 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