github.com/hellobchain/third_party@v0.0.0-20230331131523-deb0478a2e52/go-chi/chi/middleware/wrap_writer.go (about)

     1  package middleware
     2  
     3  // The original work was derived from Goji's middleware, source:
     4  // https://github.com/zenazn/goji/tree/master/web/middleware
     5  
     6  import (
     7  	"bufio"
     8  	"io"
     9  	"net"
    10  
    11  	"github.com/hellobchain/newcryptosm/http"
    12  )
    13  
    14  // WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
    15  // into various parts of the response process.
    16  type WrapResponseWriter interface {
    17  	http.ResponseWriter
    18  	// Status returns the HTTP status of the request, or 0 if one has not
    19  	// yet been sent.
    20  	Status() int
    21  	// BytesWritten returns the total number of bytes sent to the client.
    22  	BytesWritten() int
    23  	// Tee causes the response body to be written to the given io.Writer in
    24  	// addition to proxying the writes through. Only one io.Writer can be
    25  	// tee'd to at once: setting a second one will overwrite the first.
    26  	// Writes will be sent to the proxy before being written to this
    27  	// io.Writer. It is illegal for the tee'd writer to be modified
    28  	// concurrently with writes.
    29  	Tee(io.Writer)
    30  	// Unwrap returns the original proxied target.
    31  	Unwrap() http.ResponseWriter
    32  }
    33  
    34  // basicWriter wraps a http.ResponseWriter that implements the minimal
    35  // http.ResponseWriter interface.
    36  type basicWriter struct {
    37  	http.ResponseWriter
    38  	wroteHeader bool
    39  	code        int
    40  	bytes       int
    41  	tee         io.Writer
    42  }
    43  
    44  func (b *basicWriter) WriteHeader(code int) {
    45  	if !b.wroteHeader {
    46  		b.code = code
    47  		b.wroteHeader = true
    48  		b.ResponseWriter.WriteHeader(code)
    49  	}
    50  }
    51  func (b *basicWriter) Write(buf []byte) (int, error) {
    52  	b.WriteHeader(http.StatusOK)
    53  	n, err := b.ResponseWriter.Write(buf)
    54  	if b.tee != nil {
    55  		_, err2 := b.tee.Write(buf[:n])
    56  		// Prefer errors generated by the proxied writer.
    57  		if err == nil {
    58  			err = err2
    59  		}
    60  	}
    61  	b.bytes += n
    62  	return n, err
    63  }
    64  func (b *basicWriter) maybeWriteHeader() {
    65  	if !b.wroteHeader {
    66  		b.WriteHeader(http.StatusOK)
    67  	}
    68  }
    69  func (b *basicWriter) Status() int {
    70  	return b.code
    71  }
    72  func (b *basicWriter) BytesWritten() int {
    73  	return b.bytes
    74  }
    75  func (b *basicWriter) Tee(w io.Writer) {
    76  	b.tee = w
    77  }
    78  func (b *basicWriter) Unwrap() http.ResponseWriter {
    79  	return b.ResponseWriter
    80  }
    81  
    82  type flushWriter struct {
    83  	basicWriter
    84  }
    85  
    86  func (f *flushWriter) Flush() {
    87  	f.wroteHeader = true
    88  
    89  	fl := f.basicWriter.ResponseWriter.(http.Flusher)
    90  	fl.Flush()
    91  }
    92  
    93  var _ http.Flusher = &flushWriter{}
    94  
    95  // httpFancyWriter is a HTTP writer that additionally satisfies http.CloseNotifier,
    96  // http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
    97  // of wrapping the http.ResponseWriter that package http gives you, in order to
    98  // make the proxied object support the full method set of the proxied object.
    99  type httpFancyWriter struct {
   100  	basicWriter
   101  }
   102  
   103  func (f *httpFancyWriter) CloseNotify() <-chan bool {
   104  	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
   105  	return cn.CloseNotify()
   106  }
   107  func (f *httpFancyWriter) Flush() {
   108  	f.wroteHeader = true
   109  
   110  	fl := f.basicWriter.ResponseWriter.(http.Flusher)
   111  	fl.Flush()
   112  }
   113  func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
   114  	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
   115  	return hj.Hijack()
   116  }
   117  func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
   118  	if f.basicWriter.tee != nil {
   119  		n, err := io.Copy(&f.basicWriter, r)
   120  		f.basicWriter.bytes += int(n)
   121  		return n, err
   122  	}
   123  	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
   124  	f.basicWriter.maybeWriteHeader()
   125  	n, err := rf.ReadFrom(r)
   126  	f.basicWriter.bytes += int(n)
   127  	return n, err
   128  }
   129  
   130  var _ http.CloseNotifier = &httpFancyWriter{}
   131  var _ http.Flusher = &httpFancyWriter{}
   132  var _ http.Hijacker = &httpFancyWriter{}
   133  var _ io.ReaderFrom = &httpFancyWriter{}
   134  
   135  // http2FancyWriter is a HTTP2 writer that additionally satisfies http.CloseNotifier,
   136  // http.Flusher, and io.ReaderFrom. It exists for the common case
   137  // of wrapping the http.ResponseWriter that package http gives you, in order to
   138  // make the proxied object support the full method set of the proxied object.
   139  type http2FancyWriter struct {
   140  	basicWriter
   141  }
   142  
   143  func (f *http2FancyWriter) CloseNotify() <-chan bool {
   144  	cn := f.basicWriter.ResponseWriter.(http.CloseNotifier)
   145  	return cn.CloseNotify()
   146  }
   147  func (f *http2FancyWriter) Flush() {
   148  	f.wroteHeader = true
   149  
   150  	fl := f.basicWriter.ResponseWriter.(http.Flusher)
   151  	fl.Flush()
   152  }
   153  
   154  var _ http.CloseNotifier = &http2FancyWriter{}
   155  var _ http.Flusher = &http2FancyWriter{}