diff --git a/common/buf/multi_buffer.go b/common/buf/multi_buffer.go index 91b9a247..b65a1471 100644 --- a/common/buf/multi_buffer.go +++ b/common/buf/multi_buffer.go @@ -141,16 +141,17 @@ func SplitSize(mb MultiBuffer, size int32) (MultiBuffer, MultiBuffer) { totalBytes := int32(0) var r MultiBuffer - endIndex := 0 + endIndex := -1 for i := range mb { if totalBytes+mb[i].Len() > size { endIndex = i break } + totalBytes += mb[i].Len() r = append(r, mb[i]) mb[i] = nil } - if endIndex == len(mb) { + if endIndex == -1 { // To reuse mb array mb = mb[:0] } else { diff --git a/common/buf/multi_buffer_test.go b/common/buf/multi_buffer_test.go index 1e412522..8fd20d4d 100644 --- a/common/buf/multi_buffer_test.go +++ b/common/buf/multi_buffer_test.go @@ -5,6 +5,8 @@ import ( "io" "testing" + "github.com/google/go-cmp/cmp" + "v2ray.com/core/common" . "v2ray.com/core/common/buf" . "v2ray.com/ext/assert" @@ -37,13 +39,58 @@ func TestMultiBufferAppend(t *testing.T) { } func TestMultiBufferSliceBySizeLarge(t *testing.T) { - assert := With(t) - lb := make([]byte, 8*1024) common.Must2(io.ReadFull(rand.Reader, lb)) mb := MergeBytes(nil, lb) mb, mb2 := SplitSize(mb, 1024) - assert(mb2.Len(), Equals, int32(1024)) + if mb2.Len() != 1024 { + t.Error("expect length 1024, but got ", mb2.Len()) + } + if mb.Len() != 7*1024 { + t.Error("expect length 7*1024, but got ", mb.Len()) + } + + mb, mb3 := SplitSize(mb, 7*1024) + if mb3.Len() != 7*1024 { + t.Error("expect length 7*1024, but got", mb.Len()) + } + + if !mb.IsEmpty() { + t.Error("expect empty buffer, but got ", mb.Len()) + } +} + +func TestMultiBufferSplitFirst(t *testing.T) { + b1 := New() + b1.WriteString("b1") + + b2 := New() + b2.WriteString("b2") + + b3 := New() + b3.WriteString("b3") + + var mb MultiBuffer + mb = append(mb, b1, b2, b3) + + mb, c1 := SplitFirst(mb) + if diff := cmp.Diff(b1.String(), c1.String()); diff != "" { + t.Error(diff) + } + + mb, c2 := SplitFirst(mb) + if diff := cmp.Diff(b2.String(), c2.String()); diff != "" { + t.Error(diff) + } + + mb, c3 := SplitFirst(mb) + if diff := cmp.Diff(b3.String(), c3.String()); diff != "" { + t.Error(diff) + } + + if !mb.IsEmpty() { + t.Error("expect empty buffer, but got ", mb.String()) + } }