middleware/logger/tracking_response_writer.go

package logger

import (
	"bufio"
	"errors"
	"net"
	"net/http"
)

var (
	errWriterNotImplentsHijacker = errors.New("ResponseWriter does not implement the Hijacker interface")
)

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
}