From 188746fe2c3a2584c024101ff3356149e4fbf55c Mon Sep 17 00:00:00 2001
From: Alexander Tumin <iamtakingiteasy@eientei.org>
Date: Mon, 26 Jan 2026 04:48:06 +0300
Subject: [PATCH] Allow buffer reuse in ReadHeader/WriteHeader variants.

---
 read.go       | 12 ++++++++++++
 read_test.go  | 30 ++++++++++++++++++++++++++----
 write.go      | 16 ++++++++++++++--
 write_test.go | 16 ++++++++++++++++
 4 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/read.go b/read.go
index 1771816..437c088 100644
--- a/read.go
+++ b/read.go
@@ -21,6 +21,18 @@ func ReadHeader(r io.Reader) (h Header, err error) {
 	// So 14 - 2 = 12.
 	bts := make([]byte, 2, MaxHeaderSize-2)
 
+	return ReadHeaderBuffer(r, bts)
+}
+
+// ReadHeaderBuffer reads a frame header from r using user-provided buffer.
+// Provided buffer must be at least 12 bytes long.
+func ReadHeaderBuffer(r io.Reader, bts []byte) (h Header, err error) {
+	if cap(bts) < MaxHeaderSize-2 {
+		return h, io.ErrShortBuffer
+	}
+
+	bts = bts[:2]
+
 	// Prepare to hold first 2 bytes to choose size of next read.
 	_, err = io.ReadFull(r, bts)
 	if err != nil {
diff --git a/read_test.go b/read_test.go
index 7759bda..7e6bb59 100644
--- a/read_test.go
+++ b/read_test.go
@@ -38,18 +38,40 @@ func TestReadHeader(t *testing.T) {
 }
 
 func BenchmarkReadHeader(b *testing.B) {
+	setup := func(header Header, n int) (rds []io.Reader) {
+		bts := MustCompileFrame(Frame{Header: header})
+		rds = make([]io.Reader, n)
+		for i := 0; i < n; i++ {
+			rds[i] = bytes.NewReader(bts)
+		}
+
+		return
+	}
+
 	for i, bench := range RWBenchCases {
 		b.Run(fmt.Sprintf("%s#%d", bench.label, i), func(b *testing.B) {
-			bts := MustCompileFrame(Frame{Header: bench.header})
-			rds := make([]io.Reader, b.N)
+			rds := setup(bench.header, b.N)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+
 			for i := 0; i < b.N; i++ {
-				rds[i] = bytes.NewReader(bts)
+				_, err := ReadHeader(rds[i])
+				if err != nil {
+					b.Fatal(err)
+				}
 			}
+		})
+
+		b.Run(fmt.Sprintf("reused-buffer-%s#%d", bench.label, i), func(b *testing.B) {
+			rds := setup(bench.header, b.N)
+			bts := make([]byte, MaxHeaderSize)
 
+			b.ReportAllocs()
 			b.ResetTimer()
 
 			for i := 0; i < b.N; i++ {
-				_, err := ReadHeader(rds[i])
+				_, err := ReadHeaderBuffer(rds[i], bts)
 				if err != nil {
 					b.Fatal(err)
 				}
diff --git a/write.go b/write.go
index 94557c6..9e2fc07 100644
--- a/write.go
+++ b/write.go
@@ -50,11 +50,23 @@ func WriteHeader(w io.Writer, h Header) error {
 	// Make slice of bytes with capacity 14 that could hold any header.
 	bts := make([]byte, MaxHeaderSize)
 
+	return WriteHeaderBuffer(w, h, bts)
+}
+
+// WriteHeaderBuffer writes header binary representation into w using user-provided buffer.
+// Provided buffer must be at least 14 bytes long.
+func WriteHeaderBuffer(w io.Writer, h Header, bts []byte) error {
+	if cap(bts) < MaxHeaderSize {
+		return io.ErrShortBuffer
+	}
+
+	bts = bts[:MaxHeaderSize]
+
+	bts[0] = h.Rsv<<4 | byte(h.OpCode)
+
 	if h.Fin {
 		bts[0] |= bit0
 	}
-	bts[0] |= h.Rsv << 4
-	bts[0] |= byte(h.OpCode)
 
 	var n int
 	switch {
diff --git a/write_test.go b/write_test.go
index 6904ca0..e24f3c0 100644
--- a/write_test.go
+++ b/write_test.go
@@ -31,11 +31,27 @@ func TestWriteHeader(t *testing.T) {
 func BenchmarkWriteHeader(b *testing.B) {
 	for _, bench := range RWBenchCases {
 		b.Run(bench.label, func(b *testing.B) {
+			b.ReportAllocs()
+			b.ResetTimer()
+
 			for i := 0; i < b.N; i++ {
 				if err := WriteHeader(ioutil.Discard, bench.header); err != nil {
 					b.Fatal(err)
 				}
 			}
 		})
+
+		b.Run("reused-buffer-"+bench.label, func(b *testing.B) {
+			bts := make([]byte, MaxHeaderSize)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				if err := WriteHeaderBuffer(ioutil.Discard, bench.header, bts); err != nil {
+					b.Fatal(err)
+				}
+			}
+		})
 	}
 }