traefik/pkg/middlewares/accesslog/capture_response_writer.go
Jean-Baptiste Doumenjou c0f1e74bed
chore: move to Traefik organization.
Co-authored-by: Romain <rtribotte@users.noreply.github.com>
2020-09-16 15:46:04 +02:00

83 lines
1.8 KiB
Go

package accesslog
import (
"bufio"
"fmt"
"net"
"net/http"
"github.com/traefik/traefik/v2/pkg/middlewares"
)
var _ middlewares.Stateful = &captureResponseWriterWithCloseNotify{}
type capturer interface {
http.ResponseWriter
Size() int64
Status() int
}
func newCaptureResponseWriter(rw http.ResponseWriter) capturer {
capt := &captureResponseWriter{rw: rw}
if _, ok := rw.(http.CloseNotifier); !ok {
return capt
}
return &captureResponseWriterWithCloseNotify{capt}
}
// captureResponseWriter is a wrapper of type http.ResponseWriter
// that tracks request status and size.
type captureResponseWriter struct {
rw http.ResponseWriter
status int
size int64
}
type captureResponseWriterWithCloseNotify struct {
*captureResponseWriter
}
// CloseNotify returns a channel that receives at most a
// single value (true) when the client connection has gone away.
func (r *captureResponseWriterWithCloseNotify) CloseNotify() <-chan bool {
return r.rw.(http.CloseNotifier).CloseNotify()
}
func (crw *captureResponseWriter) Header() http.Header {
return crw.rw.Header()
}
func (crw *captureResponseWriter) Write(b []byte) (int, error) {
if crw.status == 0 {
crw.status = http.StatusOK
}
size, err := crw.rw.Write(b)
crw.size += int64(size)
return size, err
}
func (crw *captureResponseWriter) WriteHeader(s int) {
crw.rw.WriteHeader(s)
crw.status = s
}
func (crw *captureResponseWriter) Flush() {
if f, ok := crw.rw.(http.Flusher); ok {
f.Flush()
}
}
func (crw *captureResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if h, ok := crw.rw.(http.Hijacker); ok {
return h.Hijack()
}
return nil, nil, fmt.Errorf("not a hijacker: %T", crw.rw)
}
func (crw *captureResponseWriter) Status() int {
return crw.status
}
func (crw *captureResponseWriter) Size() int64 {
return crw.size
}