From 7f84e2e246aa9a9b5adb28ccd86fc61635d5c0ea Mon Sep 17 00:00:00 2001
From: Runxi Yu
Date: Thu, 5 Mar 2026 17:26:39 +0800
Subject: internal/zlib: Add counting flate reader

---
 internal/zlib/reader_reset.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

(limited to 'internal/zlib/reader_reset.go')

diff --git a/internal/zlib/reader_reset.go b/internal/zlib/reader_reset.go
index a39337f7..6f15b681 100644
--- a/internal/zlib/reader_reset.go
+++ b/internal/zlib/reader_reset.go
@@ -15,13 +15,17 @@ import (
 	"github.com/klauspost/compress/flate"
 )
 
-func (z *reader) Reset(r io.Reader, dict []byte) error {
-	*z = reader{decompressor: z.decompressor}
+// Reset resets receiver to read a new zlib stream.
+func (z *Reader) Reset(r io.Reader, dict []byte) error {
+	*z = Reader{decompressor: z.decompressor}
+	var input flate.Reader
 	if fr, ok := r.(flate.Reader); ok {
-		z.r = fr
+		input = fr
 	} else {
-		z.r = bufio.NewReader(r)
+		input = bufio.NewReader(r)
 	}
+	z.counter = &countingFlateReader{inner: input}
+	z.r = z.counter
 
 	// Read the header (RFC 1950 section 2.2.).
 	_, z.err = io.ReadFull(z.r, z.scratch[0:2])
-- 
cgit v1.3.1-10-gc9f91