go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/web/status_response_writer.go (about) 1 /* 2 3 Copyright (c) 2023 - Present. Will Charczuk. All rights reserved. 4 Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository. 5 6 */ 7 8 package web 9 10 import ( 11 "bufio" 12 "errors" 13 "io" 14 "net" 15 "net/http" 16 ) 17 18 var ( 19 _ ResponseWriter = (*StatusResponseWriter)(nil) 20 _ http.ResponseWriter = (*StatusResponseWriter)(nil) 21 _ http.Flusher = (*StatusResponseWriter)(nil) 22 _ io.Closer = (*StatusResponseWriter)(nil) 23 ) 24 25 // NewStatusResponseWriter creates a new response writer. 26 func NewStatusResponseWriter(w http.ResponseWriter) *StatusResponseWriter { 27 if typed, ok := w.(*StatusResponseWriter); ok { 28 return typed 29 } 30 if typed, ok := w.(ResponseWriter); ok { 31 return &StatusResponseWriter{ 32 innerResponse: typed.InnerResponse(), 33 } 34 } 35 return &StatusResponseWriter{ 36 innerResponse: w, 37 } 38 } 39 40 // StatusResponseWriter a better response writer 41 type StatusResponseWriter struct { 42 innerResponse http.ResponseWriter 43 statusCode int 44 contentLength int 45 } 46 47 // Write writes the data to the response. 48 func (rw *StatusResponseWriter) Write(b []byte) (int, error) { 49 bytesWritten, err := rw.innerResponse.Write(b) 50 rw.contentLength += bytesWritten 51 return bytesWritten, err 52 } 53 54 // Header accesses the response header collection. 55 func (rw *StatusResponseWriter) Header() http.Header { 56 return rw.innerResponse.Header() 57 } 58 59 // Hijack wraps response writer's Hijack function. 60 func (rw *StatusResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { 61 hijacker, ok := rw.innerResponse.(http.Hijacker) 62 if !ok { 63 return nil, nil, errors.New("inner responseWriter doesn't support Hijacker interface") 64 } 65 return hijacker.Hijack() 66 } 67 68 // WriteHeader writes the status code (it is a somewhat poorly chosen method name from the standard library). 69 func (rw *StatusResponseWriter) WriteHeader(code int) { 70 rw.statusCode = code 71 rw.innerResponse.WriteHeader(code) 72 } 73 74 // InnerResponse returns the backing writer. 75 func (rw *StatusResponseWriter) InnerResponse() http.ResponseWriter { 76 return rw.innerResponse 77 } 78 79 // Flush calls flush on the inner response writer if it is supported. 80 func (rw *StatusResponseWriter) Flush() { 81 if typed, ok := rw.innerResponse.(http.Flusher); ok { 82 typed.Flush() 83 } 84 } 85 86 // StatusCode returns the status code. 87 func (rw *StatusResponseWriter) StatusCode() int { 88 return rw.statusCode 89 } 90 91 // ContentLength returns the content length 92 func (rw *StatusResponseWriter) ContentLength() int { 93 return rw.contentLength 94 } 95 96 // Close calls close on the inner response if it supports it. 97 func (rw *StatusResponseWriter) Close() error { 98 if typed, ok := rw.innerResponse.(io.Closer); ok { 99 return typed.Close() 100 } 101 return nil 102 }