diff options
Diffstat (limited to 'pkg/tcpip/buffer')
-rw-r--r-- | pkg/tcpip/buffer/BUILD | 5 | ||||
-rw-r--r-- | pkg/tcpip/buffer/prependable.go | 18 | ||||
-rw-r--r-- | pkg/tcpip/buffer/prependable_test.go | 50 |
3 files changed, 66 insertions, 7 deletions
diff --git a/pkg/tcpip/buffer/BUILD b/pkg/tcpip/buffer/BUILD index 563bc78ea..5e135c50d 100644 --- a/pkg/tcpip/buffer/BUILD +++ b/pkg/tcpip/buffer/BUILD @@ -14,6 +14,9 @@ go_library( go_test( name = "buffer_test", size = "small", - srcs = ["view_test.go"], + srcs = [ + "prependable_test.go", + "view_test.go", + ], library = ":buffer", ) diff --git a/pkg/tcpip/buffer/prependable.go b/pkg/tcpip/buffer/prependable.go index ba21f4eca..57d1922ab 100644 --- a/pkg/tcpip/buffer/prependable.go +++ b/pkg/tcpip/buffer/prependable.go @@ -32,13 +32,19 @@ func NewPrependable(size int) Prependable { return Prependable{buf: NewView(size), usedIdx: size} } -// NewPrependableFromView creates an entirely-used Prependable from a View. +// NewPrependableFromView creates a Prependable from a View and allocates +// additional space if needed. // -// NewPrependableFromView takes ownership of v. Note that since the entire -// prependable is used, further attempts to call Prepend will note that size > -// p.usedIdx and return nil. -func NewPrependableFromView(v View) Prependable { - return Prependable{buf: v, usedIdx: 0} +// NewPrependableFromView takes ownership of v. Note that if the entire +// prependable is used, further attempts to call Prepend will note that +// size > p.usedIdx and return nil. +func NewPrependableFromView(v View, extraCap int) Prependable { + if extraCap == 0 { + return Prependable{buf: v, usedIdx: 0} + } + buf := make([]byte, extraCap, extraCap+len(v)) + buf = append(buf, v...) + return Prependable{buf: buf, usedIdx: extraCap} } // NewEmptyPrependableFromView creates a new prependable buffer from a View. diff --git a/pkg/tcpip/buffer/prependable_test.go b/pkg/tcpip/buffer/prependable_test.go new file mode 100644 index 000000000..435a94a61 --- /dev/null +++ b/pkg/tcpip/buffer/prependable_test.go @@ -0,0 +1,50 @@ +// Copyright 2019 The gVisor Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +package buffer + +import ( + "reflect" + "testing" +) + +func TestNewPrependableFromView(t *testing.T) { + tests := []struct { + comment string + view View + extraSize int + want Prependable + }{ + { + comment: "Reserve extra space", + view: View("abc"), + extraSize: 2, + want: Prependable{buf: View("\x00\x00abc"), usedIdx: 2}, + }, + { + comment: "Don't reserve extra space", + view: View("abc"), + extraSize: 0, + want: Prependable{buf: View("abc"), usedIdx: 0}, + }, + } + + for _, testCase := range tests { + t.Run(testCase.comment, func(t *testing.T) { + prep := NewPrependableFromView(testCase.view, testCase.extraSize) + if !reflect.DeepEqual(prep, testCase.want) { + t.Errorf("NewPrependableFromView(%#v, %d) = %#v; want %#v", testCase.view, testCase.extraSize, prep, testCase.want) + } + }) + } +} |