1
Fork 0
mirror of https://github.com/caddyserver/caddy.git synced 2024-12-16 21:56:40 -05:00
caddy/modules/caddyhttp/responsewriter.go
Matthew Holt 8bae8f5f5a
http: Always set status code via response recorder
Fixes panic if no upstream handler wrote anything to the response
2020-01-08 18:37:41 -07:00

259 lines
8.2 KiB
Go

// Copyright 2015 Matthew Holt and The Caddy Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package caddyhttp
import (
"bufio"
"bytes"
"fmt"
"io"
"net"
"net/http"
)
// ResponseWriterWrapper wraps an underlying ResponseWriter and
// promotes its Pusher/Flusher/Hijacker methods as well. To use
// this type, embed a pointer to it within your own struct type
// that implements the http.ResponseWriter interface, then call
// methods on the embedded value. You can make sure your type
// wraps correctly by asserting that it implements the
// HTTPInterfaces interface.
type ResponseWriterWrapper struct {
http.ResponseWriter
}
// Hijack implements http.Hijacker. It simply calls the underlying
// ResponseWriter's Hijack method if there is one, or returns
// ErrNotImplemented otherwise.
func (rww *ResponseWriterWrapper) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if hj, ok := rww.ResponseWriter.(http.Hijacker); ok {
return hj.Hijack()
}
return nil, nil, ErrNotImplemented
}
// Flush implements http.Flusher. It simply calls the underlying
// ResponseWriter's Flush method if there is one.
func (rww *ResponseWriterWrapper) Flush() {
if f, ok := rww.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
// Push implements http.Pusher. It simply calls the underlying
// ResponseWriter's Push method if there is one, or returns
// ErrNotImplemented otherwise.
func (rww *ResponseWriterWrapper) Push(target string, opts *http.PushOptions) error {
if pusher, ok := rww.ResponseWriter.(http.Pusher); ok {
return pusher.Push(target, opts)
}
return ErrNotImplemented
}
// HTTPInterfaces mix all the interfaces that middleware ResponseWriters need to support.
type HTTPInterfaces interface {
http.ResponseWriter
http.Pusher
http.Flusher
http.Hijacker
}
// ErrNotImplemented is returned when an underlying
// ResponseWriter does not implement the required method.
var ErrNotImplemented = fmt.Errorf("method not implemented")
type responseRecorder struct {
*ResponseWriterWrapper
wroteHeader bool
statusCode int
buf *bytes.Buffer
shouldBuffer ShouldBufferFunc
stream bool
size int
header http.Header
}
// NewResponseRecorder returns a new ResponseRecorder that can be
// used instead of a standard http.ResponseWriter. The recorder is
// useful for middlewares which need to buffer a response and
// potentially process its entire body before actually writing the
// response to the underlying writer. Of course, buffering the entire
// body has a memory overhead, but sometimes there is no way to avoid
// buffering the whole response, hence the existence of this type.
// Still, if at all practical, handlers should strive to stream
// responses by wrapping Write and WriteHeader methods instead of
// buffering whole response bodies.
//
// Buffering is actually optional. The shouldBuffer function will
// be called just before the headers are written. If it returns
// true, the headers and body will be buffered by this recorder
// and not written to the underlying writer; if false, the headers
// will be written immediately and the body will be streamed out
// directly to the underlying writer. If shouldBuffer is nil,
// the response will never be buffered and will always be streamed
// directly to the writer.
//
// You can know if shouldBuffer returned true by calling Buffered().
//
// The provided buffer buf should be obtained from a pool for best
// performance (see the sync.Pool type).
//
// Proper usage of a recorder looks like this:
//
// rec := caddyhttp.NewResponseRecorder(w, buf, shouldBuffer)
// err := next.ServeHTTP(rec, req)
// if err != nil {
// return err
// }
// if !rec.Buffered() {
// return nil
// }
// // process the buffered response here
//
// After a response has been buffered, remember that any upstream header
// manipulations are only manifest in the recorder's Header(), not the
// Header() of the underlying ResponseWriter. Thus if you wish to inspect
// or change response headers, you either need to use rec.Header(), or
// copy rec.Header() into w.Header() first (see caddyhttp.CopyHeader).
//
// Once you are ready to write the response, there are two ways you can do
// it. The easier way is to have the recorder do it:
//
// rec.WriteResponse()
//
// This writes the recorded response headers as well as the buffered body.
// Or, you may wish to do it yourself, especially if you manipulated the
// buffered body. First you will need to copy the recorded headers, then
// write the headers with the recorded status code, then write the body
// (this example writes the recorder's body buffer, but you might have
// your own body to write instead):
//
// caddyhttp.CopyHeader(w.Header(), rec.Header())
// w.WriteHeader(rec.Status())
// io.Copy(w, rec.Buffer())
//
func NewResponseRecorder(w http.ResponseWriter, buf *bytes.Buffer, shouldBuffer ShouldBufferFunc) ResponseRecorder {
// copy the current response header into this buffer so
// that any header manipulations on the buffered header
// are consistent with what would be written out
hdr := make(http.Header)
CopyHeader(hdr, w.Header())
return &responseRecorder{
ResponseWriterWrapper: &ResponseWriterWrapper{ResponseWriter: w},
buf: buf,
shouldBuffer: shouldBuffer,
header: hdr,
}
}
func (rr *responseRecorder) Header() http.Header {
return rr.header
}
func (rr *responseRecorder) WriteHeader(statusCode int) {
if rr.wroteHeader {
return
}
rr.statusCode = statusCode
rr.wroteHeader = true
// decide whether we should buffer the response
if rr.shouldBuffer == nil {
rr.stream = true
} else {
rr.stream = !rr.shouldBuffer(rr.statusCode, rr.header)
}
// if not buffered, immediately write header
if rr.stream {
CopyHeader(rr.ResponseWriterWrapper.Header(), rr.header)
rr.ResponseWriterWrapper.WriteHeader(rr.statusCode)
}
}
func (rr *responseRecorder) Write(data []byte) (int, error) {
rr.WriteHeader(http.StatusOK)
var n int
var err error
if rr.stream {
n, err = rr.ResponseWriterWrapper.Write(data)
} else {
n, err = rr.buf.Write(data)
}
if err == nil {
rr.size += n
}
return n, err
}
// Status returns the status code that was written, if any.
func (rr *responseRecorder) Status() int {
return rr.statusCode
}
// Size returns the number of bytes written,
// not including the response headers.
func (rr *responseRecorder) Size() int {
return rr.size
}
// Buffer returns the body buffer that rr was created with.
// You should still have your original pointer, though.
func (rr *responseRecorder) Buffer() *bytes.Buffer {
return rr.buf
}
// Buffered returns whether rr has decided to buffer the response.
func (rr *responseRecorder) Buffered() bool {
return !rr.stream
}
func (rr *responseRecorder) WriteResponse() error {
if rr.stream {
return nil
}
CopyHeader(rr.ResponseWriterWrapper.Header(), rr.header)
if rr.statusCode == 0 {
// could happen if no handlers actually wrote anything,
// and this prevents a panic; status must be > 0
rr.statusCode = http.StatusOK
}
rr.ResponseWriterWrapper.WriteHeader(rr.statusCode)
_, err := io.Copy(rr.ResponseWriterWrapper, rr.buf)
return err
}
// ResponseRecorder is a http.ResponseWriter that records
// responses instead of writing them to the client. See
// docs for NewResponseRecorder for proper usage.
type ResponseRecorder interface {
HTTPInterfaces
Status() int
Buffer() *bytes.Buffer
Buffered() bool
Size() int
WriteResponse() error
}
// ShouldBufferFunc is a function that returns true if the
// response should be buffered, given the pending HTTP status
// code and response headers.
type ShouldBufferFunc func(status int, header http.Header) bool
// Interface guards
var (
_ HTTPInterfaces = (*ResponseWriterWrapper)(nil)
_ ResponseRecorder = (*responseRecorder)(nil)
)