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  }