From ad737ded3cd1354cb6f4208f11fcd8f99f94c5f3 Mon Sep 17 00:00:00 2001 From: John Guo Date: Wed, 15 Feb 2023 14:13:32 +0800 Subject: [PATCH] fix issue #2447 (#2448) --- net/ghttp/ghttp_response.go | 4 +- net/ghttp/ghttp_response_write.go | 2 +- net/ghttp/ghttp_response_writer.go | 29 +++------ net/ghttp/internal/response/response.go | 8 +++ .../internal/response/response_writer.go | 59 +++++++++++++++++++ 5 files changed, 78 insertions(+), 24 deletions(-) create mode 100644 net/ghttp/internal/response/response.go create mode 100644 net/ghttp/internal/response/response_writer.go diff --git a/net/ghttp/ghttp_response.go b/net/ghttp/ghttp_response.go index 02596693e..44bb3bb7f 100644 --- a/net/ghttp/ghttp_response.go +++ b/net/ghttp/ghttp_response.go @@ -15,6 +15,7 @@ import ( "net/url" "time" + "github.com/gogf/gf/v2/net/ghttp/internal/response" "github.com/gogf/gf/v2/net/gtrace" "github.com/gogf/gf/v2/os/gfile" "github.com/gogf/gf/v2/os/gres" @@ -34,7 +35,7 @@ func newResponse(s *Server, w http.ResponseWriter) *Response { r := &Response{ Server: s, ResponseWriter: &ResponseWriter{ - writer: w, + writer: response.NewWriter(w), buffer: bytes.NewBuffer(nil), }, } @@ -152,7 +153,6 @@ func (r *Response) ClearBuffer() { // // See http.ServeContent func (r *Response) ServeContent(name string, modTime time.Time, content io.ReadSeeker) { - r.wroteHeader = true http.ServeContent(r.Writer.RawWriter(), r.Request.Request, name, modTime, content) } diff --git a/net/ghttp/ghttp_response_write.go b/net/ghttp/ghttp_response_write.go index 4ca5c6d2d..05b5c113f 100644 --- a/net/ghttp/ghttp_response_write.go +++ b/net/ghttp/ghttp_response_write.go @@ -19,7 +19,7 @@ import ( // Write writes `content` to the response buffer. func (r *Response) Write(content ...interface{}) { - if r.hijacked || len(content) == 0 { + if r.writer.IsHijacked() || len(content) == 0 { return } if r.Status == 0 { diff --git a/net/ghttp/ghttp_response_writer.go b/net/ghttp/ghttp_response_writer.go index 948cb9974..7d560d644 100644 --- a/net/ghttp/ghttp_response_writer.go +++ b/net/ghttp/ghttp_response_writer.go @@ -12,15 +12,15 @@ import ( "bytes" "net" "net/http" + + "github.com/gogf/gf/v2/net/ghttp/internal/response" ) // ResponseWriter is the custom writer for http response. type ResponseWriter struct { - Status int // HTTP status. - writer http.ResponseWriter // The underlying ResponseWriter. - buffer *bytes.Buffer // The output buffer. - hijacked bool // Mark this request is hijacked or not. - wroteHeader bool // Is header wrote or not, avoiding error: superfluous/multiple response.WriteHeader call. + Status int // HTTP status. + writer *response.Writer // The underlying ResponseWriter. + buffer *bytes.Buffer // The output buffer. } // RawWriter returns the underlying ResponseWriter. @@ -46,18 +46,16 @@ func (w *ResponseWriter) WriteHeader(status int) { // Hijack implements the interface function of http.Hijacker.Hijack. func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { - w.hijacked = true - return w.writer.(http.Hijacker).Hijack() + return w.writer.Hijack() } // Flush outputs the buffer to clients and clears the buffer. func (w *ResponseWriter) Flush() { - if w.hijacked { + if w.writer.IsHijacked() { return } - if w.Status != 0 && !w.isHeaderWritten() { - w.wroteHeader = true + if w.Status != 0 && !w.writer.IsHeaderWrote() { w.writer.WriteHeader(w.Status) } // Default status text output. @@ -69,14 +67,3 @@ func (w *ResponseWriter) Flush() { w.buffer.Reset() } } - -// isHeaderWrote checks and returns whether the header is written. -func (w *ResponseWriter) isHeaderWritten() bool { - if w.wroteHeader { - return true - } - if _, ok := w.writer.Header()[responseHeaderContentLength]; ok { - return true - } - return false -} diff --git a/net/ghttp/internal/response/response.go b/net/ghttp/internal/response/response.go new file mode 100644 index 000000000..71afabac4 --- /dev/null +++ b/net/ghttp/internal/response/response.go @@ -0,0 +1,8 @@ +// Copyright GoFrame Author(https://goframe.org). All Rights Reserved. +// +// This Source Code Form is subject to the terms of the MIT License. +// If a copy of the MIT was not distributed with this file, +// You can obtain one at https://github.com/gogf/gf. + +// Package response provides wrapper for http.response. +package response diff --git a/net/ghttp/internal/response/response_writer.go b/net/ghttp/internal/response/response_writer.go new file mode 100644 index 000000000..b974b6c34 --- /dev/null +++ b/net/ghttp/internal/response/response_writer.go @@ -0,0 +1,59 @@ +// Copyright GoFrame Author(https://goframe.org). All Rights Reserved. +// +// This Source Code Form is subject to the terms of the MIT License. +// If a copy of the MIT was not distributed with this file, +// You can obtain one at https://github.com/gogf/gf. + +package response + +import ( + "bufio" + "net" + "net/http" +) + +// Writer wraps http.ResponseWriter for extra features. +type Writer struct { + http.ResponseWriter // The underlying ResponseWriter. + hijacked bool // Mark this request is hijacked or not. + wroteHeader bool // Is header wrote or not, avoiding error: superfluous/multiple response.WriteHeader call. +} + +// NewWriter creates and returns a new Writer. +func NewWriter(writer http.ResponseWriter) *Writer { + return &Writer{ + ResponseWriter: writer, + } +} + +// WriteHeader implements the interface of http.ResponseWriter.WriteHeader. +func (w *Writer) WriteHeader(status int) { + w.ResponseWriter.WriteHeader(status) + w.wroteHeader = true +} + +// Hijack implements the interface function of http.Hijacker.Hijack. +func (w *Writer) Hijack() (conn net.Conn, writer *bufio.ReadWriter, err error) { + conn, writer, err = w.ResponseWriter.(http.Hijacker).Hijack() + w.hijacked = true + return +} + +// IsHeaderWrote returns if the header status is written. +func (w *Writer) IsHeaderWrote() bool { + return w.wroteHeader +} + +// IsHijacked returns if the connection is hijacked. +func (w *Writer) IsHijacked() bool { + return w.hijacked +} + +// Flush sends any buffered data to the client. +func (w *Writer) Flush() { + flusher, ok := w.ResponseWriter.(http.Flusher) + if ok { + flusher.Flush() + w.wroteHeader = true + } +}