From b836026204c5aa46f1c91049f7d76ce7ce72e261 Mon Sep 17 00:00:00 2001
From: lesismal <lesismal@gmail.com>
Date: Tue, 23 Apr 2024 18:07:56 +0800
Subject: [PATCH] http: fix body reader append logic

---
 conn_unix.go        |  1 +
 nbhttp/body.go      |  7 ++++++-
 nbhttp/body_test.go | 10 +++++++---
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/conn_unix.go b/conn_unix.go
index e2861582..f00e8f19 100644
--- 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
index 03ba5858..700536a7 100644
--- 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
index 8144fe67..3cd6539e 100644
--- 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)