github.com/msales/pkg/v3@v3.24.0/httpx/middleware/response_writer.go (about)

     1  package middleware
     2  
     3  import (
     4  	"net/http"
     5  )
     6  
     7  // ResponseWriter is a wrapper around http.ResponseWriter that provides extra
     8  // information about the response.
     9  type ResponseWriter interface {
    10  	http.ResponseWriter
    11  
    12  	// Status returns the status code of the response or 0 if the response has
    13  	// not be written.
    14  	Status() int
    15  }
    16  
    17  // NewResponseWriter create a new ResponseWriter.
    18  func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
    19  	return &responseWriter{rw, 0}
    20  }
    21  
    22  type responseWriter struct {
    23  	http.ResponseWriter
    24  	status int
    25  }
    26  
    27  // Write writes the data to the connection as part of an HTTP reply.
    28  func (rw *responseWriter) Write(b []byte) (int, error) {
    29  	if rw.status == 0 {
    30  		rw.status = http.StatusOK
    31  	}
    32  
    33  	return rw.ResponseWriter.Write(b)
    34  }
    35  
    36  // WriteHeader sends an HTTP response header with status code.
    37  func (rw *responseWriter) WriteHeader(s int) {
    38  	rw.status = s
    39  	rw.ResponseWriter.WriteHeader(s)
    40  }
    41  
    42  // Status returns the status code of the response or 0 if the response has
    43  // not be written.
    44  func (rw *responseWriter) Status() int {
    45  	return rw.status
    46  }