fix prompting superfluous header (#20)

sillygod · Jun 24, 2020 · 5800ff1 · 5800ff1
1 parent b19f2e3
commit 5800ff1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/cache.go b/cache.go
@@ -318,9 +318,12 @@ func (e *Entry) writePublicResponse(w http.ResponseWriter) error {
 }
 
 func (e *Entry) writePrivateResponse(w http.ResponseWriter) error {
+	// wrap the original response writer
 	e.Response.SetBody(backends.WrapResponseWriterToBackend(w))
 	e.Response.WaitClose()
-	w.WriteHeader(e.Response.Code)
+	if !e.Response.IsFirstByteWritten {
+		w.WriteHeader(e.Response.Code)
+	}
 	return nil
 }
 

diff --git a/response.go b/response.go
@@ -26,8 +26,9 @@ type Response struct {
 	body       backends.Backend
 	snapHeader http.Header
 
-	wroteHeader  bool
-	bodyComplete bool
+	wroteHeader        bool
+	bodyComplete       bool
+	IsFirstByteWritten bool
 
 	bodyChan         chan struct{} // indicates whether the backend is set or not.
 	bodyCompleteChan chan struct{}
@@ -94,6 +95,9 @@ func (r *Response) Write(buf []byte) (int, error) {
 	}
 
 	if r.body != nil {
+		if !r.IsFirstByteWritten {
+			r.IsFirstByteWritten = true
+		}
 		return r.body.Write(buf)
 	}