PixivFE/server/middleware/logger.go

// server/middleware/logger.go
package middleware

import (
	"net/http"
	"strings"
	"time"

	"codeberg.org/vnpower/pixivfe/v2/audit"
	"codeberg.org/vnpower/pixivfe/v2/config"
	"codeberg.org/vnpower/pixivfe/v2/server/request_context"
)

// ResponseWriterInterceptStatus wraps http.ResponseWriter to intercept the status code.
type ResponseWriterInterceptStatus struct {
	statusCode int
	http.ResponseWriter
}

// WriteHeader intercepts the status code before writing it to the response.
func (w *ResponseWriterInterceptStatus) WriteHeader(code int) {
	w.statusCode = code
	w.ResponseWriter.WriteHeader(code)
}

// CanRequestSkipLogger determines if a request should bypass the logging middleware.
// This is useful for reducing log clutter from static assets and development-specific routes.
func CanRequestSkipLogger(r *http.Request) bool {
	// NOTE: Uncomment the following line to log all requests
	// return false
	path := r.URL.Path
	return strings.HasPrefix(path, "/img/") ||
		strings.HasPrefix(path, "/css/") ||
		strings.HasPrefix(path, "/js/") ||
		strings.HasPrefix(path, "/diagnostics") ||
		(config.GlobalConfig.InDevelopment &&
			(strings.HasPrefix(path, "/proxy/s.pximg.net/") || strings.HasPrefix(path, "/proxy/i.pximg.net/")))
}

// LogRequest is a middleware that logs incoming HTTP requests and their corresponding responses.
func LogRequest(auditor *audit.Auditor) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if CanRequestSkipLogger(r) {
				// If skipping, simply call the next handler without logging
				next.ServeHTTP(w, r)
				return
			}

			// Wrap the ResponseWriter to capture the status code
			wrappedWriter := &ResponseWriterInterceptStatus{
				ResponseWriter: w,
			}

			startTime := time.Now()

			next.ServeHTTP(wrappedWriter, r)

			endTime := time.Now()

			ctx := request_context.Get(r)

			auditor.LogServerRoundTrip(audit.ServerRequestSpan{
				StartTime:  startTime,
				EndTime:    endTime,
				RequestId:  ctx.RequestId,
				Method:     r.Method,
				Path:       r.URL.Path,
				Status:     wrappedWriter.statusCode,
				Referer:    r.Referer(),
				RemoteAddr: r.RemoteAddr,
				ErrorField: ctx.CaughtError,
			})
		})
	}
}