github.com/blend/go-sdk@v1.20220411.3/webutil/status_response_writer.go (about)

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