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

     1  package response
     2  
     3  import (
     4  	"bufio"
     5  	"errors"
     6  	"net"
     7  	"net/http"
     8  )
     9  
    10  // HeaderOnceResponseWriter is a http.ResponseWriter implementation
    11  // that only the first WriterHeader takes effect.
    12  type HeaderOnceResponseWriter struct {
    13  	w           http.ResponseWriter
    14  	wroteHeader bool
    15  }
    16  
    17  // NewHeaderOnceResponseWriter returns a HeaderOnceResponseWriter.
    18  func NewHeaderOnceResponseWriter(w http.ResponseWriter) http.ResponseWriter {
    19  	return &HeaderOnceResponseWriter{w: w}
    20  }
    21  
    22  // Flush flushes the response writer.
    23  func (w *HeaderOnceResponseWriter) Flush() {
    24  	if flusher, ok := w.w.(http.Flusher); ok {
    25  		flusher.Flush()
    26  	}
    27  }
    28  
    29  // Header returns the http header.
    30  func (w *HeaderOnceResponseWriter) Header() http.Header {
    31  	return w.w.Header()
    32  }
    33  
    34  // Hijack implements the http.Hijacker interface.
    35  // This expands the Response to fulfill http.Hijacker if the underlying http.ResponseWriter supports it.
    36  func (w *HeaderOnceResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    37  	if hijacked, ok := w.w.(http.Hijacker); ok {
    38  		return hijacked.Hijack()
    39  	}
    40  
    41  	return nil, nil, errors.New("server doesn't support hijacking")
    42  }
    43  
    44  // Write writes bytes into w.
    45  func (w *HeaderOnceResponseWriter) Write(bytes []byte) (int, error) {
    46  	return w.w.Write(bytes)
    47  }
    48  
    49  // WriteHeader writes code into w, and not sealing the writer.
    50  func (w *HeaderOnceResponseWriter) WriteHeader(code int) {
    51  	if w.wroteHeader {
    52  		return
    53  	}
    54  
    55  	w.w.WriteHeader(code)
    56  	w.wroteHeader = true
    57  }