From c6a051a5a646666cfccad23f0a6a450f55077929 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Sat, 7 Oct 2023 03:07:46 +0800
Subject: [PATCH] 1

---
 web/Web.go | 35 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/web/Web.go b/web/Web.go
index 79aced2..d7b6b2b 100644
--- a/web/Web.go
+++ b/web/Web.go
@@ -349,36 +349,35 @@ func (t *Cache) Cache(key string, aliveDur time.Duration, w http.ResponseWriter)
 }
 
 type withflush struct {
-	f  func()
-	h  func() http.Header
-	w  func([]byte) (int, error)
-	wh func(int)
+	raw http.ResponseWriter
 }
 
 func (t withflush) Header() http.Header {
-	return t.h()
+	if t.raw != nil {
+		return t.raw.Header()
+	}
+	return make(http.Header)
 }
 func (t withflush) Write(b []byte) (i int, e error) {
-	i, e = t.w(b)
-	if t.f != nil {
-		t.f()
+	if t.raw != nil {
+		i, e = t.Write(b)
+		if e != nil {
+			return
+		}
+		if flusher, ok := t.raw.(http.Flusher); ok {
+			flusher.Flush()
+		}
 	}
 	return
 }
 func (t withflush) WriteHeader(i int) {
-	t.wh(i)
+	if t.raw != nil {
+		t.raw.WriteHeader(i)
+	}
 }
 
 func WithFlush(w http.ResponseWriter) http.ResponseWriter {
-	if Flusher, ok := w.(http.Flusher); ok {
-		return withflush{
-			f:  Flusher.Flush,
-			h:  w.Header,
-			w:  w.Write,
-			wh: w.WriteHeader,
-		}
-	}
-	return w
+	return withflush{w}
 }
 
 func WithStatusCode(w http.ResponseWriter, code int) {
-- 
2.39.2