lazyhandler/middleware/httplog/log.go

package httplog

import (
	"bufio"
	"context"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"time"

	"git.jeffthecoder.xyz/public/lazyhandler/magic"
)

type ctxKey int

const (
	loggerKey ctxKey = iota
)

var (
	ErrNotHijackable = errors.New("not a hijacker")
)

type Log struct {
	LogStart      bool
	LogStartLevel slog.Level

	LogFinish      bool
	LogFinishLevel slog.Level
}

type responseRecorder struct {
	http.ResponseWriter
	StatusCode int
}

func (recorder *responseRecorder) WriteHeader(statusCode int) {
	recorder.StatusCode = statusCode

	recorder.ResponseWriter.WriteHeader(statusCode)
}

func (recorder *responseRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	if hijacker, ok := recorder.ResponseWriter.(http.Hijacker); ok {
		recorder.ResponseWriter = nil
		return hijacker.Hijack()
	}

	return nil, nil, ErrNotHijackable
}

var (
	_ http.Hijacker = &responseRecorder{}
)

func (log Log) WrapHandler(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		args := []any{
			"remote_addr", r.RemoteAddr,
			"host", r.Host,
			"path", r.URL.Path,
		}

		startTime := time.Now()
		if log.LogStart {
			slog.With(append(args, "time", startTime)...).Log(r.Context(), log.LogStartLevel, "request")
		}
		recorder := &responseRecorder{ResponseWriter: w, StatusCode: 200}
		next.ServeHTTP(recorder, r.WithContext(context.WithValue(r.Context(), loggerKey, slog.With(args...))))
		if log.LogFinish && recorder.ResponseWriter != nil {
			finishTime := time.Now()
			slog.With(append(args, "time", finishTime, "duration", finishTime.Sub(startTime), "status", recorder.StatusCode)...).Log(r.Context(), log.LogFinishLevel, "response")
		}
	})
}

func Logger(r *http.Request) *slog.Logger {
	if logger, ok := r.Context().Value(loggerKey).(*slog.Logger); ok {
		return logger.With("time", time.Now())
	}

	return slog.With(
		"remote_addr", r.RemoteAddr,
		"host", r.Host,
		"path", r.URL.Path,
		"time", time.Now(),
	)
}

func RegisterExtractor() {
	magic.RegisterExtractorGeneric[*slog.Logger](func(r *http.Request) (any, error) {
		return Logger(r), nil
	})
}