http: fix body reader append logic

lesismal · Apr 23, 2024 · b836026 · b836026
1 parent 64090fd
commit b836026
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/conn_unix.go b/conn_unix.go
@@ -57,6 +57,7 @@ func (c *Conn) newToWriteBuf(buf []byte) {
 			if cap(tail.buf) < tailLen+l {
 				b := c.p.g.BodyAllocator.Malloc(tailLen + l)[:tailLen]
 				copy(b, tail.buf)
+				c.p.g.BodyAllocator.Free(tail.buf)
 				tail.buf = b
 			}
 			tail.buf = append(tail.buf, buf...)

diff --git a/nbhttp/body.go b/nbhttp/body.go
@@ -38,6 +38,7 @@ func (br *BodyReader) Read(p []byte) (int, error) {
 		nc := copy(p[ncopy:], b[br.index:])
 		if nc+br.index >= len(b) {
 			br.engine.BodyAllocator.Free(b)
+			br.buffers[0] = nil
 			br.buffers = br.buffers[1:]
 			br.index = 0
 		} else {
@@ -118,7 +119,11 @@ func (br *BodyReader) append(data []byte) error {
 		l := len(b)
 		bLeft := cap(b) - len(b)
 		if bLeft > 0 {
-			b = b[:cap(b)]
+			if bLeft > len(data) {
+				b = b[:l+len(data)]
+			} else {
+				b = b[:cap(b)]
+			}
 			nc := copy(b[l:], data)
 			data = data[nc:]
 			br.buffers[i] = b

diff --git a/nbhttp/body_test.go b/nbhttp/body_test.go
@@ -5,13 +5,17 @@ import (
 	"crypto/rand"
 	"io"
 	"testing"
+
+	"github.com/lesismal/nbio/mempool"
 )
 
 func TestBodyReader(t *testing.T) {
-	engine := NewEngine(Config{})
-	b1 := make([]byte, 0, 1997)
+	engine := NewEngine(Config{
+		BodyAllocator: mempool.NewAligned(),
+	})
+	b1 := make([]byte, 2049)
 	rand.Read(b1)
-	b2 := make([]byte, 4096)
+	b2 := make([]byte, 1132)
 	rand.Read(b2)
 	b3 := make([]byte, 11111)
 	rand.Read(b3)