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 }