54 lines
1.2 KiB
Go
54 lines
1.2 KiB
Go
|
package logger
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bufio"
|
||
|
|
"errors"
|
||
|
|
"net"
|
||
|
|
"net/http"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
errWriterNotImplentsHijacker = errors.New("ResponseWriter does not implement the Hijacker interface") // nolint:golint
|
||
|
|
)
|
||
|
|
|
||
|
|
type trackingResponseWriter struct {
|
||
|
|
http.ResponseWriter
|
||
|
|
|
||
|
|
size int
|
||
|
|
status int
|
||
|
|
}
|
||
|
|
|
||
|
|
func newTrackingResponseWriter(w http.ResponseWriter) *trackingResponseWriter {
|
||
|
|
return &trackingResponseWriter{ResponseWriter: w, status: 200}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WriteHeader implements http.ResponseWriter and saves status
|
||
|
|
func (s *trackingResponseWriter) WriteHeader(status int) {
|
||
|
|
s.status = status
|
||
|
|
s.ResponseWriter.WriteHeader(status)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Write implements http.ResponseWriter and tracks number of bytes written
|
||
|
|
func (s *trackingResponseWriter) Write(b []byte) (int, error) {
|
||
|
|
size, err := s.ResponseWriter.Write(b)
|
||
|
|
s.size += size
|
||
|
|
|
||
|
|
return size, err
|
||
|
|
}
|
||
|
|
|
||
|
|
// Flush implements http.Flusher
|
||
|
|
func (s *trackingResponseWriter) Flush() {
|
||
|
|
if f, ok := s.ResponseWriter.(http.Flusher); ok {
|
||
|
|
f.Flush()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Hijack implements http.Hijacker
|
||
|
|
func (s *trackingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
||
|
|
if hj, ok := s.ResponseWriter.(http.Hijacker); ok {
|
||
|
|
return hj.Hijack()
|
||
|
|
}
|
||
|
|
|
||
|
|
return nil, nil, errWriterNotImplentsHijacker
|
||
|
|
}
|