go-zero/rest/internal/security/withcoderesponsewriter.go
Kevin Wan e97e1f10db
simplify code with http.Flusher type conversion (#325)
* simplify code with http.Flusher type conversion

* simplify code with http.Flusher type conversion, better version
2020-12-29 10:25:55 +08:00

28 lines
521 B
Go

package security
import "net/http"
type WithCodeResponseWriter struct {
Writer http.ResponseWriter
Code int
}
func (w *WithCodeResponseWriter) Flush() {
if flusher, ok := w.Writer.(http.Flusher); ok {
flusher.Flush()
}
}
func (w *WithCodeResponseWriter) Header() http.Header {
return w.Writer.Header()
}
func (w *WithCodeResponseWriter) Write(bytes []byte) (int, error) {
return w.Writer.Write(bytes)
}
func (w *WithCodeResponseWriter) WriteHeader(code int) {
w.Writer.WriteHeader(code)
w.Code = code
}