soft-serve

 1package web
 2
 3import (
 4	"bufio"
 5	"fmt"
 6	"net"
 7	"net/http"
 8	"time"
 9
10	"github.com/charmbracelet/log/v2"
11	"github.com/dustin/go-humanize"
12)
13
14// logWriter is a wrapper around http.ResponseWriter that allows us to capture
15// the HTTP status code and bytes written to the response.
16type logWriter struct {
17	http.ResponseWriter
18	code, bytes int
19}
20
21var _ http.ResponseWriter = (*logWriter)(nil)
22
23var _ http.Flusher = (*logWriter)(nil)
24
25var _ http.Hijacker = (*logWriter)(nil)
26
27var _ http.CloseNotifier = (*logWriter)(nil) // nolint: staticcheck
28
29// Write implements http.ResponseWriter.
30func (r *logWriter) Write(p []byte) (int, error) {
31	written, err := r.ResponseWriter.Write(p)
32	r.bytes += written
33	return written, err
34}
35
36// Note this is generally only called when sending an HTTP error, so it's
37// important to set the `code` value to 200 as a default.
38func (r *logWriter) WriteHeader(code int) {
39	r.code = code
40	r.ResponseWriter.WriteHeader(code)
41}
42
43// Unwrap returns the underlying http.ResponseWriter.
44func (r *logWriter) Unwrap() http.ResponseWriter {
45	return r.ResponseWriter
46}
47
48// Flush implements http.Flusher.
49func (r *logWriter) Flush() {
50	if f, ok := r.ResponseWriter.(http.Flusher); ok {
51		f.Flush()
52	}
53}
54
55// CloseNotify implements http.CloseNotifier.
56func (r *logWriter) CloseNotify() <-chan bool {
57	if cn, ok := r.ResponseWriter.(http.CloseNotifier); ok { // nolint: staticcheck
58		return cn.CloseNotify()
59	}
60	return nil
61}
62
63// Hijack implements http.Hijacker.
64func (r *logWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
65	if h, ok := r.ResponseWriter.(http.Hijacker); ok {
66		return h.Hijack()
67	}
68	return nil, nil, fmt.Errorf("http.Hijacker not implemented")
69}
70
71// NewLoggingMiddleware returns a new logging middleware.
72func NewLoggingMiddleware(next http.Handler, logger *log.Logger) http.Handler {
73	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
74		start := time.Now()
75		writer := &logWriter{code: http.StatusOK, ResponseWriter: w}
76		logger.Debug("request",
77			"method", r.Method,
78			"path", r.URL,
79			"addr", r.RemoteAddr)
80		next.ServeHTTP(writer, r)
81		elapsed := time.Since(start)
82		logger.Debug("response",
83			"status", fmt.Sprintf("%d %s", writer.code, http.StatusText(writer.code)),
84			"bytes", humanize.Bytes(uint64(writer.bytes)), //nolint:gosec
85			"time", elapsed)
86	})
87}