github.com/blend/go-sdk@v1.20220411.3/datadog/traceserver/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 traceserver
     9  
    10  import (
    11  	"bufio"
    12  	"fmt"
    13  	"io"
    14  	"net"
    15  	"net/http"
    16  )
    17  
    18  // ResponseWriter a better response writer
    19  type ResponseWriter struct {
    20  	http.ResponseWriter
    21  
    22  	StatusCode    int
    23  	ContentLength int
    24  }
    25  
    26  // Write writes the data to the response.
    27  func (rw *ResponseWriter) Write(b []byte) (int, error) {
    28  	bytesWritten, err := rw.ResponseWriter.Write(b)
    29  	rw.ContentLength += bytesWritten
    30  	return bytesWritten, err
    31  }
    32  
    33  // Hijack wraps response writer's Hijack function.
    34  func (rw *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    35  	hijacker, ok := rw.ResponseWriter.(http.Hijacker)
    36  	if !ok {
    37  		return nil, nil, fmt.Errorf("response writer hijack; wrapped response writer doesn't support the hijacker interface")
    38  	}
    39  	return hijacker.Hijack()
    40  }
    41  
    42  // WriteHeader writes the status code (it is a somewhat poorly chosen method name from the standard library).
    43  func (rw *ResponseWriter) WriteHeader(code int) {
    44  	rw.StatusCode = code
    45  	rw.ResponseWriter.WriteHeader(code)
    46  }
    47  
    48  // Flush calls flush on the inner response writer if it is supported.
    49  func (rw *ResponseWriter) Flush() {
    50  	if typed, ok := rw.ResponseWriter.(http.Flusher); ok {
    51  		typed.Flush()
    52  	}
    53  }
    54  
    55  // Close calls close on the inner response if it supports it.
    56  func (rw *ResponseWriter) Close() error {
    57  	if typed, ok := rw.ResponseWriter.(io.Closer); ok {
    58  		return typed.Close()
    59  	}
    60  	return nil
    61  }