You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/rest/internal/response/withcoderesponsewriter.go

48 lines
1.2 KiB
Go

package response
import (
"bufio"
"errors"
"net"
"net/http"
)
// A WithCodeResponseWriter is a helper to delay sealing a http.ResponseWriter on writing code.
type WithCodeResponseWriter struct {
Writer http.ResponseWriter
Code int
}
// Flush flushes the response writer.
func (w *WithCodeResponseWriter) Flush() {
if flusher, ok := w.Writer.(http.Flusher); ok {
flusher.Flush()
}
}
// Header returns the http header.
func (w *WithCodeResponseWriter) Header() http.Header {
return w.Writer.Header()
}
// Hijack implements the http.Hijacker interface.
// This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hijacked, ok := w.Writer.(http.Hijacker); ok {
return hijacked.Hijack()
}
return nil, nil, errors.New("server doesn't support hijacking")
}
// Write writes bytes into w.
func (w *WithCodeResponseWriter) Write(bytes []byte) (int, error) {
return w.Writer.Write(bytes)
}
// WriteHeader writes code into w, and not sealing the writer.
func (w *WithCodeResponseWriter) WriteHeader(code int) {
w.Writer.WriteHeader(code)
w.Code = code
}