github.com/searKing/golang/go@v1.2.117/net/http/response_writer.go (about)

     1  // Copyright 2023 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package http
     6  
     7  import (
     8  	"bufio"
     9  	"context"
    10  	"errors"
    11  	"fmt"
    12  	"io"
    13  	"net"
    14  	"net/http"
    15  )
    16  
    17  var (
    18  	_ http.ResponseWriter = (*recordResponseWriter)(nil)
    19  	_ http.ResponseWriter = (*recordResponseWriter)(nil)
    20  	_ http.Hijacker       = (*recordResponseWriter)(nil)
    21  	_ http.Flusher        = (*recordResponseWriter)(nil)
    22  	_ http.CloseNotifier  = (*recordResponseWriter)(nil)
    23  )
    24  
    25  // NewRecordResponseWriter creates a ResponseWriter that is a wrapper around [http.ResponseWriter] that
    26  // provides extra information about the response.
    27  // It is recommended that middleware handlers use this construct to wrap a [http.ResponseWriter]
    28  // if the functionality calls for it.
    29  func NewRecordResponseWriter(rw http.ResponseWriter) *recordResponseWriter {
    30  	return &recordResponseWriter{
    31  		ResponseWriter: rw,
    32  	}
    33  }
    34  
    35  type recordResponseWriter struct {
    36  	http.ResponseWriter
    37  	size   int
    38  	status int
    39  }
    40  
    41  func (w *recordResponseWriter) Context() context.Context {
    42  	return w.Context()
    43  }
    44  
    45  func (w *recordResponseWriter) Unwrap() http.ResponseWriter {
    46  	return w.ResponseWriter
    47  }
    48  
    49  func (w *recordResponseWriter) reset(writer http.ResponseWriter) {
    50  	w.ResponseWriter = writer
    51  	w.size = 0
    52  	w.status = 0
    53  }
    54  
    55  func (w *recordResponseWriter) WriteHeader(code int) {
    56  	w.status = code
    57  	w.ResponseWriter.WriteHeader(code)
    58  }
    59  
    60  // ExplicitlyWriteHeader forces to write the http header (status code + headers).
    61  // If WriteHeader is not called explicitly, the first call to Write
    62  // will trigger an implicit WriteHeader(http.StatusOK).
    63  // Thus explicit calls to WriteHeader are mainly used to
    64  // send error codes or 1xx informational responses.
    65  func (w *recordResponseWriter) ExplicitlyWriteHeader() {
    66  	if !w.Written() {
    67  		// The status will be StatusOK if WriteHeader has not been called yet
    68  		w.WriteHeader(http.StatusOK)
    69  	}
    70  }
    71  
    72  func (w *recordResponseWriter) Write(data []byte) (n int, err error) {
    73  	w.ExplicitlyWriteHeader()
    74  	n, err = w.ResponseWriter.Write(data)
    75  	w.size += n
    76  	return
    77  }
    78  
    79  // WriteString writes the string into the response body.
    80  func (w *recordResponseWriter) WriteString(s string) (n int, err error) {
    81  	w.ExplicitlyWriteHeader()
    82  	n, err = io.WriteString(w.ResponseWriter, s)
    83  	w.size += n
    84  	return
    85  }
    86  
    87  // Status returns the status code of the response or 0 if the response has
    88  // not been written
    89  func (w *recordResponseWriter) Status() int {
    90  	return w.status
    91  }
    92  
    93  // Size returns the number of bytes already written into the response http body.
    94  // See Written()
    95  func (w *recordResponseWriter) Size() int {
    96  	return w.size
    97  }
    98  
    99  // Written returns whether or not the ResponseWriter has been written.
   100  func (w *recordResponseWriter) Written() bool {
   101  	return w.status != 0
   102  }
   103  
   104  // Hijack implements the http.Hijacker interface.
   105  func (w *recordResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
   106  	if w.size < 0 {
   107  		w.size = 0
   108  	}
   109  	if hijacker, ok := w.ResponseWriter.(http.Hijacker); ok {
   110  		return hijacker.Hijack()
   111  
   112  	}
   113  	return nil, nil, errors.New("ResponseWriter doesn't support Hijacker interface")
   114  }
   115  
   116  // CloseNotify implements the http.CloseNotifier interface.
   117  func (w *recordResponseWriter) CloseNotify() <-chan bool {
   118  	if gone, ok := w.ResponseWriter.(http.CloseNotifier); ok {
   119  		return gone.CloseNotify()
   120  	}
   121  	return nil
   122  }
   123  
   124  // Flush implements the http.Flusher interface.
   125  func (w *recordResponseWriter) Flush() {
   126  	if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
   127  		w.ExplicitlyWriteHeader()
   128  		flusher.Flush()
   129  	}
   130  }
   131  
   132  // Pusher get the http.Pusher for server push
   133  func (w *recordResponseWriter) Pusher() (pusher http.Pusher) {
   134  	if pusher, ok := w.ResponseWriter.(http.Pusher); ok {
   135  		return pusher
   136  	}
   137  	return nil
   138  }
   139  
   140  func (w *recordResponseWriter) Push(target string, opts *http.PushOptions) error {
   141  	pusher, ok := w.ResponseWriter.(http.Pusher)
   142  	if ok {
   143  		return pusher.Push(target, opts)
   144  	}
   145  	return fmt.Errorf("the ResponseWriter doesn't support the Pusher interface")
   146  }