github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/middleware/response.go (about)

     1  package middleware
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"net/http"
     9  )
    10  
    11  const (
    12  	maxResponseBodyInLogs = 4096 // At most 4k bytes from response bodies in our logs.
    13  )
    14  
    15  type badResponseLoggingWriter interface {
    16  	http.ResponseWriter
    17  	getStatusCode() int
    18  	getWriteError() error
    19  }
    20  
    21  // nonFlushingBadResponseLoggingWriter writes the body of "bad" responses (i.e. 5xx
    22  // responses) to a buffer.
    23  type nonFlushingBadResponseLoggingWriter struct {
    24  	rw            http.ResponseWriter
    25  	buffer        io.Writer
    26  	logBody       bool
    27  	bodyBytesLeft int
    28  	statusCode    int
    29  	writeError    error // The error returned when downstream Write() fails.
    30  }
    31  
    32  // flushingBadResponseLoggingWriter is a badResponseLoggingWriter that
    33  // implements http.Flusher.
    34  type flushingBadResponseLoggingWriter struct {
    35  	nonFlushingBadResponseLoggingWriter
    36  	f http.Flusher
    37  }
    38  
    39  func newBadResponseLoggingWriter(rw http.ResponseWriter, buffer io.Writer) badResponseLoggingWriter {
    40  	b := nonFlushingBadResponseLoggingWriter{
    41  		rw:            rw,
    42  		buffer:        buffer,
    43  		logBody:       false,
    44  		bodyBytesLeft: maxResponseBodyInLogs,
    45  		statusCode:    http.StatusOK,
    46  	}
    47  
    48  	if f, ok := rw.(http.Flusher); ok {
    49  		return &flushingBadResponseLoggingWriter{b, f}
    50  	}
    51  
    52  	return &b
    53  }
    54  
    55  // Unwrap method is used by http.ResponseController to get access to original http.ResponseWriter.
    56  func (b *nonFlushingBadResponseLoggingWriter) Unwrap() http.ResponseWriter {
    57  	return b.rw
    58  }
    59  
    60  // Header returns the header map that will be sent by WriteHeader.
    61  // Implements ResponseWriter.
    62  func (b *nonFlushingBadResponseLoggingWriter) Header() http.Header {
    63  	return b.rw.Header()
    64  }
    65  
    66  // Write writes HTTP response data.
    67  func (b *nonFlushingBadResponseLoggingWriter) Write(data []byte) (int, error) {
    68  	if b.statusCode == 0 {
    69  		// WriteHeader has (probably) not been called, so we need to call it with StatusOK to fulfill the interface contract.
    70  		// https://godoc.org/net/http#ResponseWriter
    71  		b.WriteHeader(http.StatusOK)
    72  	}
    73  	n, err := b.rw.Write(data)
    74  	if b.logBody {
    75  		b.captureResponseBody(data)
    76  	}
    77  	if err != nil {
    78  		b.writeError = err
    79  	}
    80  	return n, err
    81  }
    82  
    83  // WriteHeader writes the HTTP response header.
    84  func (b *nonFlushingBadResponseLoggingWriter) WriteHeader(statusCode int) {
    85  	b.statusCode = statusCode
    86  	if statusCode >= 500 {
    87  		b.logBody = true
    88  	}
    89  	b.rw.WriteHeader(statusCode)
    90  }
    91  
    92  // Hijack hijacks the first response writer that is a Hijacker.
    93  func (b *nonFlushingBadResponseLoggingWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    94  	hj, ok := b.rw.(http.Hijacker)
    95  	if ok {
    96  		return hj.Hijack()
    97  	}
    98  	return nil, nil, fmt.Errorf("badResponseLoggingWriter: can't cast underlying response writer to Hijacker")
    99  }
   100  
   101  func (b *nonFlushingBadResponseLoggingWriter) getStatusCode() int {
   102  	return b.statusCode
   103  }
   104  
   105  func (b *nonFlushingBadResponseLoggingWriter) getWriteError() error {
   106  	return b.writeError
   107  }
   108  
   109  func (b *nonFlushingBadResponseLoggingWriter) captureResponseBody(data []byte) {
   110  	if len(data) > b.bodyBytesLeft {
   111  		b.buffer.Write(data[:b.bodyBytesLeft])
   112  		io.WriteString(b.buffer, "...")
   113  		b.bodyBytesLeft = 0
   114  		b.logBody = false
   115  	} else {
   116  		b.buffer.Write(data)
   117  		b.bodyBytesLeft -= len(data)
   118  	}
   119  }
   120  
   121  func (b *flushingBadResponseLoggingWriter) Flush() {
   122  	b.f.Flush()
   123  }