Bläddra i källkod

limit size of written data

Darien Raymond 9 år sedan
förälder
incheckning
b575de2a55
2 ändrade filer med 18 tillägg och 2 borttagningar
  1. 5 2
      common/alloc/buffer.go
  2. 13 0
      common/alloc/buffer_test.go

+ 5 - 2
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().

+ 13 - 0
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()