log/access_logger.go

package log

import (
	"bufio"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/sirupsen/logrus"
	"gitlab.com/gitlab-org/labkit/correlation"
	"gitlab.com/gitlab-org/labkit/mask"
	"gitlab.com/gitlab-org/labkit/remoteip"
)

// AccessLogger will generate access logs for the service.
func AccessLogger(h http.Handler, opts ...AccessLoggerOption) http.Handler {
	config := applyAccessLoggerOptions(opts)

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var bodyLength atomic.Int64

		r.Body = &bodyByteCounter{
			body:  r.Body,
			count: &bodyLength,
		}

		lrw := newLoggingResponseWriter(w, &config, &bodyLength)
		defer lrw.requestFinished(r)

		h.ServeHTTP(lrw, r)
		lrw.setStatus()
	})
}

func newLoggingResponseWriter(rw http.ResponseWriter, config *accessLoggerConfig, bodyByteCount *atomic.Int64) notifiableResponseWriter {
	out := loggingResponseWriter{
		rw:            rw,
		started:       time.Now(),
		config:        config,
		bodyByteCount: bodyByteCount,
	}

	// If the underlying response writer supports hijacking,
	// we need to ensure that we do too
	if _, ok := rw.(http.Hijacker); ok {
		return &hijackingResponseWriter{out}
	}

	return &out
}

// notifiableResponseWriter is a response writer that can be notified when the request is complete,
// via the requestFinished method.
type notifiableResponseWriter interface {
	http.ResponseWriter

	// requestFinished is called by the middleware when the request has completed
	requestFinished(r *http.Request)
	setStatus()
}

type loggingResponseWriter struct {
	rw            http.ResponseWriter
	status        int
	wroteHeader   bool
	written       int64
	started       time.Time
	ttfb          time.Duration
	config        *accessLoggerConfig
	contentType   string
	bodyByteCount *atomic.Int64
}

func (l *loggingResponseWriter) Header() http.Header {
	return l.rw.Header()
}

func (l *loggingResponseWriter) Write(data []byte) (int, error) {
	if !l.wroteHeader {
		l.WriteHeader(http.StatusOK)
	}

	n, err := l.rw.Write(data)

	l.written += int64(n)

	return n, err
}

func (l *loggingResponseWriter) WriteHeader(status int) {
	if l.wroteHeader {
		return
	}

	l.wroteHeader = true
	l.status = status
	l.contentType = l.Header().Get("Content-Type")
	l.ttfb = time.Since(l.started)

	l.rw.WriteHeader(status)
}

// Unwrap lets http.ResponseController get the underlying http.ResponseWriter,
// by implementing the [rwUnwrapper](https://cs.opensource.google/go/go/+/refs/tags/go1.21.4:src/net/http/responsecontroller.go;l=42-44) interface.
func (l *loggingResponseWriter) Unwrap() http.ResponseWriter {
	return l.rw
}

func (l *loggingResponseWriter) setStatus() {
	if !l.wroteHeader {
		// If upstream never called WriteHeader, the Go net/http server will
		// respond with status 200 to the client. We should also log status 200
		// in that case.
		l.status = http.StatusOK
	}
}

//nolint:cyclop
func (l *loggingResponseWriter) accessLogFields(r *http.Request) logrus.Fields {
	duration := time.Since(l.started)

	fields := l.config.extraFields(r)
	fieldsBitMask := l.config.fields

	// Optionally add built in fields
	if fieldsBitMask&CorrelationID != 0 {
		fields[correlation.FieldName] = correlation.ExtractFromContext(r.Context())
	}

	if fieldsBitMask&HTTPHost != 0 {
		fields[httpHostField] = r.Host
	}

	if fieldsBitMask&HTTPRemoteIP != 0 {
		fields[httpRemoteIPField] = l.getRemoteIP(r)
	}

	if fieldsBitMask&HTTPRemoteAddr != 0 {
		fields[httpRemoteAddrField] = r.RemoteAddr
	}

	if fieldsBitMask&HTTPRequestMethod != 0 {
		fields[httpRequestMethodField] = r.Method
	}

	if fieldsBitMask&HTTPURI != 0 {
		fields[httpURIField] = mask.URL(r.RequestURI)
	}

	if fieldsBitMask&HTTPProto != 0 {
		fields[httpProtoField] = r.Proto
	}

	if fieldsBitMask&HTTPResponseStatusCode != 0 {
		fields[httpResponseStatusCodeField] = l.status
	}

	if fieldsBitMask&HTTPResponseSize != 0 {
		fields[httpResponseSizeField] = l.written
	}

	if fieldsBitMask&HTTPRequestSize != 0 {
		fields[httpRequestSizeField] = computeApproximateRequestSize(r, l.bodyByteCount.Load())
	}

	if fieldsBitMask&HTTPRequestReferrer != 0 {
		fields[httpRequestReferrerField] = mask.URL(r.Referer())
	}

	if fieldsBitMask&HTTPUserAgent != 0 {
		fields[httpUserAgentField] = r.UserAgent()
	}

	if fieldsBitMask&RequestDuration != 0 {
		fields[requestDurationField] = int64(duration / time.Millisecond)
	}

	if fieldsBitMask&RequestTTFB != 0 && l.ttfb > 0 {
		fields[requestTTFBField] = l.ttfb.Milliseconds()
	}

	if fieldsBitMask&System != 0 {
		fields[systemField] = "http"
	}

	if fieldsBitMask&HTTPResponseContentType != 0 {
		fields[httpResponseContentTypeField] = l.contentType
	}

	return fields
}

func (l *loggingResponseWriter) requestFinished(r *http.Request) {
	l.config.logger.WithFields(l.accessLogFields(r)).Info("access")
}

func (l *loggingResponseWriter) getRemoteIP(r *http.Request) string {
	sip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil || sip == "" {
		return ""
	}

	if !l.config.xffAllowed(sip) {
		return sip
	}

	parser := remoteip.NewRemoteIPParser(l.config.trustedProxies)
	remote_ip, err := parser.GetIP(r)
	if err == nil {
		return remote_ip
	}

	return sip
}

// hijackingResponseWriter is like a loggingResponseWriter that supports the http.Hijacker interface.
type hijackingResponseWriter struct {
	loggingResponseWriter
}

func (l *hijackingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	// The only way to get here is through NewStatsCollectingResponseWriter(), which
	// checks that this cast will be valid.
	hijacker := l.rw.(http.Hijacker)
	return hijacker.Hijack()
}