gitlab.com/gitlab-org/labkit@v1.21.0/correlation/inbound_http_options.go

package correlation

import (
	"net"
	"net/http"

	"github.com/sebest/xff"
	"github.com/sirupsen/logrus"
)

// XFFAllowedFunc decides whether X-Forwarded-For headers are to be trusted.
type XFFAllowedFunc func(ip string) bool

// The configuration for InjectCorrelationID.
type inboundHandlerConfig struct {
	propagation                  bool
	sendResponseHeader           bool
	invalidCIDRsForPropagation   bool
	trustedCIDRsForPropagation   []net.IPNet
	trustedCIDRsForXForwardedFor []net.IPNet
	xffAllowed                   XFFAllowedFunc
}

// InboundHandlerOption will configure a correlation handler
// currently there are no options, but this gives us the option
// to extend the interface in a backwards compatible way.
type InboundHandlerOption func(*inboundHandlerConfig)

func applyInboundHandlerOptions(opts []InboundHandlerOption) inboundHandlerConfig {
	config := inboundHandlerConfig{
		propagation: false,
	}
	for _, v := range opts {
		v(&config)
	}

	return config
}

// WithPropagation will configure the handler to propagate existing correlation_ids
// passed in from upstream services.
// This is not the default behaviour.
func WithPropagation() InboundHandlerOption {
	return func(config *inboundHandlerConfig) {
		config.propagation = true
	}
}

// WithSetResponseHeader will configure the handler to set the correlation_id
// in the http response headers.
func WithSetResponseHeader() InboundHandlerOption {
	return func(config *inboundHandlerConfig) {
		config.sendResponseHeader = true
	}
}

// WithCIDRsTrustedForPropagation will configure the handler to set a list of trusted
// CIDR blocks to allow propagation of correlation_ids.
func WithCIDRsTrustedForPropagation(trustedCIDRs []string) InboundHandlerOption {
	return func(config *inboundHandlerConfig) {
		for _, s := range trustedCIDRs {
			_, ipNet, err := net.ParseCIDR(s)
			if err != nil {
				logrus.Errorf("Bad trusted CIDR for propagation %s: %v, propagation disabled", s, err)
				config.invalidCIDRsForPropagation = true
			} else {
				config.trustedCIDRsForPropagation = append(config.trustedCIDRsForPropagation, *ipNet)
			}
		}
	}
}

// WithCIDRsTrustedForXForwardedFor will configure the handler to trust
// X-Forwarded-For from trusted CIDR blocks.
func WithCIDRsTrustedForXForwardedFor(trustedCIDRs []string) InboundHandlerOption {
	return func(config *inboundHandlerConfig) {
		for _, s := range trustedCIDRs {
			_, ipNet, err := net.ParseCIDR(s)
			if err != nil {
				logrus.Errorf("Bad trusted CIDR for XForwardedFor %s: %v", s, err)
			} else {
				config.trustedCIDRsForXForwardedFor = append(config.trustedCIDRsForXForwardedFor, *ipNet)
			}
		}

		config.xffAllowed = func(ip string) bool {
			return isTrustedIP(ip, config.trustedCIDRsForXForwardedFor)
		}
	}
}

func isTrustedIP(ipAddress string, trustedCIDRs []net.IPNet) bool {
	ip := net.ParseIP(ipAddress)

	for _, cidr := range trustedCIDRs {
		if cidr.Contains(ip) {
			return true
		}
	}

	return false
}

func isTrustedAddr(addr string, trustedCIDRs []net.IPNet) bool {
	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		return false
	}

	return isTrustedIP(host, trustedCIDRs)
}

func (c *inboundHandlerConfig) shouldPropagate(r *http.Request) bool {
	if !c.propagation || c.invalidCIDRsForPropagation {
		return false
	}

	if len(c.trustedCIDRsForPropagation) == 0 {
		return true
	}

	remoteAddr := r.RemoteAddr

	if c.xffAllowed != nil {
		// Unix domain sockets have a remote addr of @. This will make the
		// xff package lookup the X-Forwarded-For address if available.
		if remoteAddr == "@" {
			r.RemoteAddr = "127.0.0.1:0"
		}

		remoteAddr = xff.GetRemoteAddrIfAllowed(r, c.xffAllowed)

		// If X-Forwarded-For was allowed and returned a different address, check
		// the original address against our trusted CIDRs for propagation, in case
		// the reverse proxy is trusted
		if remoteAddr != r.RemoteAddr && isTrustedAddr(r.RemoteAddr, c.trustedCIDRsForPropagation) {
			return true
		}
	}

	return isTrustedAddr(remoteAddr, c.trustedCIDRsForPropagation)
}