recorder.go

     1  // Copyright 2011 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package httptest
     6  
     7  import (
     8  	"bytes"
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"net/textproto"
    13  	"strconv"
    14  	"strings"
    15  
    16  	"golang.org/x/net/http/httpguts"
    17  )
    18  
    19  // ResponseRecorder is an implementation of [http.ResponseWriter] that
    20  // records its mutations for later inspection in tests.
    21  type ResponseRecorder struct {
    22  	// Code is the HTTP response code set by WriteHeader.
    23  	//
    24  	// Note that if a Handler never calls WriteHeader or Write,
    25  	// this might end up being 0, rather than the implicit
    26  	// http.StatusOK. To get the implicit value, use the Result
    27  	// method.
    28  	Code int
    29  
    30  	// HeaderMap contains the headers explicitly set by the Handler.
    31  	// It is an internal detail.
    32  	//
    33  	// Deprecated: HeaderMap exists for historical compatibility
    34  	// and should not be used. To access the headers returned by a handler,
    35  	// use the Response.Header map as returned by the Result method.
    36  	HeaderMap http.Header
    37  
    38  	// Body is the buffer to which the Handler's Write calls are sent.
    39  	// If nil, the Writes are silently discarded.
    40  	Body *bytes.Buffer
    41  
    42  	// Flushed is whether the Handler called Flush.
    43  	Flushed bool
    44  
    45  	result      *http.Response // cache of Result's return value
    46  	snapHeader  http.Header    // snapshot of HeaderMap at first Write
    47  	wroteHeader bool
    48  }
    49  
    50  // NewRecorder returns an initialized [ResponseRecorder].
    51  func NewRecorder() *ResponseRecorder {
    52  	return &ResponseRecorder{
    53  		HeaderMap: make(http.Header),
    54  		Body:      new(bytes.Buffer),
    55  		Code:      200,
    56  	}
    57  }
    58  
    59  // DefaultRemoteAddr is the default remote address to return in RemoteAddr if
    60  // an explicit DefaultRemoteAddr isn't set on [ResponseRecorder].
    61  const DefaultRemoteAddr = "1.2.3.4"
    62  
    63  // Header implements [http.ResponseWriter]. It returns the response
    64  // headers to mutate within a handler. To test the headers that were
    65  // written after a handler completes, use the [ResponseRecorder.Result] method and see
    66  // the returned Response value's Header.
    67  func (rw *ResponseRecorder) Header() http.Header {
    68  	m := rw.HeaderMap
    69  	if m == nil {
    70  		m = make(http.Header)
    71  		rw.HeaderMap = m
    72  	}
    73  	return m
    74  }
    75  
    76  // writeHeader writes a header if it was not written yet and
    77  // detects Content-Type if needed.
    78  //
    79  // bytes or str are the beginning of the response body.
    80  // We pass both to avoid unnecessarily generate garbage
    81  // in rw.WriteString which was created for performance reasons.
    82  // Non-nil bytes win.
    83  func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
    84  	if rw.wroteHeader {
    85  		return
    86  	}
    87  	if len(str) > 512 {
    88  		str = str[:512]
    89  	}
    90  
    91  	m := rw.Header()
    92  
    93  	_, hasType := m["Content-Type"]
    94  	hasTE := m.Get("Transfer-Encoding") != ""
    95  	if !hasType && !hasTE {
    96  		if b == nil {
    97  			b = []byte(str)
    98  		}
    99  		m.Set("Content-Type", http.DetectContentType(b))
   100  	}
   101  
   102  	rw.WriteHeader(200)
   103  }
   104  
   105  // Write implements http.ResponseWriter. The data in buf is written to
   106  // rw.Body, if not nil.
   107  func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
   108  	code := rw.Code
   109  	if !bodyAllowedForStatus(code) {
   110  		return 0, http.ErrBodyNotAllowed
   111  	}
   112  	rw.writeHeader(buf, "")
   113  	if rw.Body != nil {
   114  		rw.Body.Write(buf)
   115  	}
   116  	return len(buf), nil
   117  }
   118  
   119  // WriteString implements [io.StringWriter]. The data in str is written
   120  // to rw.Body, if not nil.
   121  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
   122  	code := rw.Code
   123  	if !bodyAllowedForStatus(code) {
   124  		return 0, http.ErrBodyNotAllowed
   125  	}
   126  	rw.writeHeader(nil, str)
   127  	if rw.Body != nil {
   128  		rw.Body.WriteString(str)
   129  	}
   130  	return len(str), nil
   131  }
   132  
   133  // bodyAllowedForStatus reports whether a given response status code
   134  // permits a body. See RFC 7230, section 3.3.
   135  func bodyAllowedForStatus(status int) bool {
   136  	switch {
   137  	case status >= 100 && status <= 199:
   138  		return false
   139  	case status == 204:
   140  		return false
   141  	case status == 304:
   142  		return false
   143  	}
   144  	return true
   145  }
   146  
   147  func checkWriteHeaderCode(code int) {
   148  	// Issue 22880: require valid WriteHeader status codes.
   149  	// For now we only enforce that it's three digits.
   150  	// In the future we might block things over 599 (600 and above aren't defined
   151  	// at https://httpwg.org/specs/rfc7231.html#status.codes)
   152  	// and we might block under 200 (once we have more mature 1xx support).
   153  	// But for now any three digits.
   154  	//
   155  	// We used to send "HTTP/1.1 000 0" on the wire in responses but there's
   156  	// no equivalent bogus thing we can realistically send in HTTP/2,
   157  	// so we'll consistently panic instead and help people find their bugs
   158  	// early. (We can't return an error from WriteHeader even if we wanted to.)
   159  	if code < 100 || code > 999 {
   160  		panic(fmt.Sprintf("invalid WriteHeader code %v", code))
   161  	}
   162  }
   163  
   164  // WriteHeader implements [http.ResponseWriter].
   165  func (rw *ResponseRecorder) WriteHeader(code int) {
   166  	if rw.wroteHeader {
   167  		return
   168  	}
   169  
   170  	checkWriteHeaderCode(code)
   171  	rw.Code = code
   172  	rw.wroteHeader = true
   173  	if rw.HeaderMap == nil {
   174  		rw.HeaderMap = make(http.Header)
   175  	}
   176  	rw.snapHeader = rw.HeaderMap.Clone()
   177  }
   178  
   179  // Flush implements [http.Flusher]. To test whether Flush was
   180  // called, see rw.Flushed.
   181  func (rw *ResponseRecorder) Flush() {
   182  	if !rw.wroteHeader {
   183  		rw.WriteHeader(200)
   184  	}
   185  	rw.Flushed = true
   186  }
   187  
   188  // Result returns the response generated by the handler.
   189  //
   190  // The returned Response will have at least its StatusCode,
   191  // Header, Body, and optionally Trailer populated.
   192  // More fields may be populated in the future, so callers should
   193  // not DeepEqual the result in tests.
   194  //
   195  // The Response.Header is a snapshot of the headers at the time of the
   196  // first write call, or at the time of this call, if the handler never
   197  // did a write.
   198  //
   199  // The Response.Body is guaranteed to be non-nil and Body.Read call is
   200  // guaranteed to not return any error other than [io.EOF].
   201  //
   202  // Result must only be called after the handler has finished running.
   203  func (rw *ResponseRecorder) Result() *http.Response {
   204  	if rw.result != nil {
   205  		return rw.result
   206  	}
   207  	if rw.snapHeader == nil {
   208  		rw.snapHeader = rw.HeaderMap.Clone()
   209  	}
   210  	res := &http.Response{
   211  		Proto:      "HTTP/1.1",
   212  		ProtoMajor: 1,
   213  		ProtoMinor: 1,
   214  		StatusCode: rw.Code,
   215  		Header:     rw.snapHeader,
   216  	}
   217  	rw.result = res
   218  	if res.StatusCode == 0 {
   219  		res.StatusCode = 200
   220  	}
   221  	res.Status = fmt.Sprintf("%03d %s", res.StatusCode, http.StatusText(res.StatusCode))
   222  	if rw.Body != nil {
   223  		res.Body = io.NopCloser(bytes.NewReader(rw.Body.Bytes()))
   224  	} else {
   225  		res.Body = http.NoBody
   226  	}
   227  	res.ContentLength = parseContentLength(res.Header.Get("Content-Length"))
   228  
   229  	if trailers, ok := rw.snapHeader["Trailer"]; ok {
   230  		res.Trailer = make(http.Header, len(trailers))
   231  		for _, k := range trailers {
   232  			for k := range strings.SplitSeq(k, ",") {
   233  				k = http.CanonicalHeaderKey(textproto.TrimString(k))
   234  				if !httpguts.ValidTrailerHeader(k) {
   235  					// Ignore since forbidden by RFC 7230, section 4.1.2.
   236  					continue
   237  				}
   238  				vv, ok := rw.HeaderMap[k]
   239  				if !ok {
   240  					continue
   241  				}
   242  				vv2 := make([]string, len(vv))
   243  				copy(vv2, vv)
   244  				res.Trailer[k] = vv2
   245  			}
   246  		}
   247  	}
   248  	for k, vv := range rw.HeaderMap {
   249  		if !strings.HasPrefix(k, http.TrailerPrefix) {
   250  			continue
   251  		}
   252  		if res.Trailer == nil {
   253  			res.Trailer = make(http.Header)
   254  		}
   255  		for _, v := range vv {
   256  			res.Trailer.Add(strings.TrimPrefix(k, http.TrailerPrefix), v)
   257  		}
   258  	}
   259  	return res
   260  }
   261  
   262  // parseContentLength trims whitespace from s and returns -1 if no value
   263  // is set, or the value if it's >= 0.
   264  //
   265  // This a modified version of same function found in net/http/transfer.go. This
   266  // one just ignores an invalid header.
   267  func parseContentLength(cl string) int64 {
   268  	cl = textproto.TrimString(cl)
   269  	if cl == "" {
   270  		return -1
   271  	}
   272  	n, err := strconv.ParseUint(cl, 10, 63)
   273  	if err != nil {
   274  		return -1
   275  	}
   276  	return int64(n)
   277  }
   278
View as plain text