github.com/lingyao2333/mo-zero@v1.4.1/rest/internal/response/withcoderesponsewriter.go (about)

     1  package response
     2  
     3  import (
     4  	"bufio"
     5  	"errors"
     6  	"net"
     7  	"net/http"
     8  )
     9  
    10  // A WithCodeResponseWriter is a helper to delay sealing a http.ResponseWriter on writing code.
    11  type WithCodeResponseWriter struct {
    12  	Writer http.ResponseWriter
    13  	Code   int
    14  }
    15  
    16  // Flush flushes the response writer.
    17  func (w *WithCodeResponseWriter) Flush() {
    18  	if flusher, ok := w.Writer.(http.Flusher); ok {
    19  		flusher.Flush()
    20  	}
    21  }
    22  
    23  // Header returns the http header.
    24  func (w *WithCodeResponseWriter) Header() http.Header {
    25  	return w.Writer.Header()
    26  }
    27  
    28  // Hijack implements the http.Hijacker interface.
    29  // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
    30  func (w *WithCodeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    31  	if hijacked, ok := w.Writer.(http.Hijacker); ok {
    32  		return hijacked.Hijack()
    33  	}
    34  
    35  	return nil, nil, errors.New("server doesn't support hijacking")
    36  }
    37  
    38  // Write writes bytes into w.
    39  func (w *WithCodeResponseWriter) Write(bytes []byte) (int, error) {
    40  	return w.Writer.Write(bytes)
    41  }
    42  
    43  // WriteHeader writes code into w, and not sealing the writer.
    44  func (w *WithCodeResponseWriter) WriteHeader(code int) {
    45  	w.Writer.WriteHeader(code)
    46  	w.Code = code
    47  }