recorder.go

Documentation: net/http/httptest

		 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 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  	rw.writeHeader(buf, "")
	 109  	if rw.Body != nil {
	 110  		rw.Body.Write(buf)
	 111  	}
	 112  	return len(buf), nil
	 113  }
	 114  
	 115  // WriteString implements io.StringWriter. The data in str is written
	 116  // to rw.Body, if not nil.
	 117  func (rw *ResponseRecorder) WriteString(str string) (int, error) {
	 118  	rw.writeHeader(nil, str)
	 119  	if rw.Body != nil {
	 120  		rw.Body.WriteString(str)
	 121  	}
	 122  	return len(str), nil
	 123  }
	 124  
	 125  func checkWriteHeaderCode(code int) {
	 126  	// Issue 22880: require valid WriteHeader status codes.
	 127  	// For now we only enforce that it's three digits.
	 128  	// In the future we might block things over 599 (600 and above aren't defined
	 129  	// at https://httpwg.org/specs/rfc7231.html#status.codes)
	 130  	// and we might block under 200 (once we have more mature 1xx support).
	 131  	// But for now any three digits.
	 132  	//
	 133  	// We used to send "HTTP/1.1 000 0" on the wire in responses but there's
	 134  	// no equivalent bogus thing we can realistically send in HTTP/2,
	 135  	// so we'll consistently panic instead and help people find their bugs
	 136  	// early. (We can't return an error from WriteHeader even if we wanted to.)
	 137  	if code < 100 || code > 999 {
	 138  		panic(fmt.Sprintf("invalid WriteHeader code %v", code))
	 139  	}
	 140  }
	 141  
	 142  // WriteHeader implements http.ResponseWriter.
	 143  func (rw *ResponseRecorder) WriteHeader(code int) {
	 144  	if rw.wroteHeader {
	 145  		return
	 146  	}
	 147  
	 148  	checkWriteHeaderCode(code)
	 149  	rw.Code = code
	 150  	rw.wroteHeader = true
	 151  	if rw.HeaderMap == nil {
	 152  		rw.HeaderMap = make(http.Header)
	 153  	}
	 154  	rw.snapHeader = rw.HeaderMap.Clone()
	 155  }
	 156  
	 157  // Flush implements http.Flusher. To test whether Flush was
	 158  // called, see rw.Flushed.
	 159  func (rw *ResponseRecorder) Flush() {
	 160  	if !rw.wroteHeader {
	 161  		rw.WriteHeader(200)
	 162  	}
	 163  	rw.Flushed = true
	 164  }
	 165  
	 166  // Result returns the response generated by the handler.
	 167  //
	 168  // The returned Response will have at least its StatusCode,
	 169  // Header, Body, and optionally Trailer populated.
	 170  // More fields may be populated in the future, so callers should
	 171  // not DeepEqual the result in tests.
	 172  //
	 173  // The Response.Header is a snapshot of the headers at the time of the
	 174  // first write call, or at the time of this call, if the handler never
	 175  // did a write.
	 176  //
	 177  // The Response.Body is guaranteed to be non-nil and Body.Read call is
	 178  // guaranteed to not return any error other than io.EOF.
	 179  //
	 180  // Result must only be called after the handler has finished running.
	 181  func (rw *ResponseRecorder) Result() *http.Response {
	 182  	if rw.result != nil {
	 183  		return rw.result
	 184  	}
	 185  	if rw.snapHeader == nil {
	 186  		rw.snapHeader = rw.HeaderMap.Clone()
	 187  	}
	 188  	res := &http.Response{
	 189  		Proto:			"HTTP/1.1",
	 190  		ProtoMajor: 1,
	 191  		ProtoMinor: 1,
	 192  		StatusCode: rw.Code,
	 193  		Header:		 rw.snapHeader,
	 194  	}
	 195  	rw.result = res
	 196  	if res.StatusCode == 0 {
	 197  		res.StatusCode = 200
	 198  	}
	 199  	res.Status = fmt.Sprintf("%03d %s", res.StatusCode, http.StatusText(res.StatusCode))
	 200  	if rw.Body != nil {
	 201  		res.Body = io.NopCloser(bytes.NewReader(rw.Body.Bytes()))
	 202  	} else {
	 203  		res.Body = http.NoBody
	 204  	}
	 205  	res.ContentLength = parseContentLength(res.Header.Get("Content-Length"))
	 206  
	 207  	if trailers, ok := rw.snapHeader["Trailer"]; ok {
	 208  		res.Trailer = make(http.Header, len(trailers))
	 209  		for _, k := range trailers {
	 210  			k = http.CanonicalHeaderKey(k)
	 211  			if !httpguts.ValidTrailerHeader(k) {
	 212  				// Ignore since forbidden by RFC 7230, section 4.1.2.
	 213  				continue
	 214  			}
	 215  			vv, ok := rw.HeaderMap[k]
	 216  			if !ok {
	 217  				continue
	 218  			}
	 219  			vv2 := make([]string, len(vv))
	 220  			copy(vv2, vv)
	 221  			res.Trailer[k] = vv2
	 222  		}
	 223  	}
	 224  	for k, vv := range rw.HeaderMap {
	 225  		if !strings.HasPrefix(k, http.TrailerPrefix) {
	 226  			continue
	 227  		}
	 228  		if res.Trailer == nil {
	 229  			res.Trailer = make(http.Header)
	 230  		}
	 231  		for _, v := range vv {
	 232  			res.Trailer.Add(strings.TrimPrefix(k, http.TrailerPrefix), v)
	 233  		}
	 234  	}
	 235  	return res
	 236  }
	 237  
	 238  // parseContentLength trims whitespace from s and returns -1 if no value
	 239  // is set, or the value if it's >= 0.
	 240  //
	 241  // This a modified version of same function found in net/http/transfer.go. This
	 242  // one just ignores an invalid header.
	 243  func parseContentLength(cl string) int64 {
	 244  	cl = textproto.TrimString(cl)
	 245  	if cl == "" {
	 246  		return -1
	 247  	}
	 248  	n, err := strconv.ParseUint(cl, 10, 63)
	 249  	if err != nil {
	 250  		return -1
	 251  	}
	 252  	return int64(n)
	 253  }
	 254
View as plain text