Source file src/github.com/henvic/httpretty/recorder.go

Documentation: github.com/henvic/httpretty

     1  package httpretty
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"net/http"
     7  )
     8  
     9  type bodyCloser struct {
    10  	r     io.Reader
    11  	close func() error
    12  }
    13  
    14  func (bc *bodyCloser) Read(p []byte) (n int, err error) {
    15  	return bc.r.Read(p)
    16  }
    17  
    18  func (bc *bodyCloser) Close() error {
    19  	return bc.close()
    20  }
    21  
    22  func newBodyReaderBuf(buf io.Reader, body io.ReadCloser) *bodyCloser {
    23  	return &bodyCloser{
    24  		r:     io.MultiReader(buf, body),
    25  		close: body.Close,
    26  	}
    27  }
    28  
    29  type responseRecorder struct {
    30  	http.ResponseWriter
    31  
    32  	statusCode int
    33  
    34  	maxReadableBody int64
    35  	size            int64
    36  	buf             *bytes.Buffer
    37  }
    38  
    39  // Write the data to the connection as part of an HTTP reply, and records it.
    40  func (rr *responseRecorder) Write(p []byte) (int, error) {
    41  	rr.size += int64(len(p))
    42  
    43  	if rr.maxReadableBody > 0 && rr.size > rr.maxReadableBody {
    44  		rr.buf = nil
    45  		return rr.ResponseWriter.Write(p)
    46  	}
    47  
    48  	defer rr.buf.Write(p)
    49  	return rr.ResponseWriter.Write(p)
    50  }
    51  
    52  // WriteHeader sends an HTTP response header with the provided
    53  // status code, and records it.
    54  func (rr *responseRecorder) WriteHeader(statusCode int) {
    55  	rr.ResponseWriter.WriteHeader(statusCode)
    56  	rr.statusCode = statusCode
    57  }
    58  

View as plain text