From b575de2a55a30a4e1511f79f3797f6da8dcd905b Mon Sep 17 00:00:00 2001 From: Darien Raymond Date: Fri, 2 Dec 2016 14:35:28 +0100 Subject: [PATCH] limit size of written data --- common/alloc/buffer.go | 7 +++++-- common/alloc/buffer_test.go | 13 +++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/common/alloc/buffer.go b/common/alloc/buffer.go index bbc359ed9..ec04189d5 100644 --- a/common/alloc/buffer.go +++ b/common/alloc/buffer.go @@ -167,8 +167,11 @@ func (b *Buffer) IsFull() bool { // Write implements Write method in io.Writer. func (b *Buffer) Write(data []byte) (int, error) { - b.Append(data) - return len(data), nil + begin := b.Len() + b.Value = b.Value[:cap(b.Value)] + nBytes := copy(b.Value[begin:], data) + b.Value = b.Value[:begin+nBytes] + return nBytes, nil } // Read implements io.Reader.Read(). diff --git a/common/alloc/buffer_test.go b/common/alloc/buffer_test.go index 4d3ed2787..4785c8174 100644 --- a/common/alloc/buffer_test.go +++ b/common/alloc/buffer_test.go @@ -59,6 +59,19 @@ func TestBufferString(t *testing.T) { assert.String(buffer.String()).Equals("Test String") } +func TestBufferWrite(t *testing.T) { + assert := assert.On(t) + + buffer := NewLocalBuffer(24).Clear() // 16 + 8 + nBytes, err := buffer.Write([]byte("abcd")) + assert.Error(err).IsNil() + assert.Int(nBytes).Equals(4) + nBytes, err = buffer.Write([]byte("abcde")) + assert.Error(err).IsNil() + assert.Int(nBytes).Equals(4) + assert.String(buffer.String()).Equals("abcdabcd") +} + func BenchmarkNewBuffer8192(b *testing.B) { for i := 0; i < b.N; i++ { buffer := NewBuffer()