proxy-pool/internal/gateway/server/protection.go

package server

import (
	"context"
	"crypto/subtle"
	"encoding/base64"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"strconv"
	"strings"

	"github.com/proxy-pool/proxy-pool/internal/gateway/policy"
)

type HTTPError struct {
	StatusCode int
	Header     http.Header
	Cause      error
}

func (err *HTTPError) Error() string {
	if err.Cause != nil {
		return err.Cause.Error()
	}
	return http.StatusText(err.StatusCode)
}

func (err *HTTPError) Unwrap() error { return err.Cause }

type BasicAuthGuard struct {
	username string
	password string
}

func NewBasicAuthGuard(username, password string) *BasicAuthGuard {
	return &BasicAuthGuard{username: username, password: password}
}

func (guard *BasicAuthGuard) Check(_ context.Context, request *http.Request) error {
	username, password, ok := parseBasicCredentials(request.Header.Get("Proxy-Authorization"))
	if ok && constantTimeEqual(username, guard.username) && constantTimeEqual(password, guard.password) {
		return nil
	}
	return &HTTPError{
		StatusCode: http.StatusProxyAuthRequired,
		Header:     http.Header{"Proxy-Authenticate": []string{`Basic realm="proxy"`}},
		Cause:      errors.New("proxy authentication failed"),
	}
}

type APIKeyGuard struct {
	header string
	value  string
}

func NewAPIKeyGuard(header, value string) *APIKeyGuard {
	if strings.TrimSpace(header) == "" {
		header = "X-API-Key"
	}
	return &APIKeyGuard{header: header, value: value}
}

func (guard *APIKeyGuard) Check(_ context.Context, request *http.Request) error {
	if constantTimeEqual(request.Header.Get(guard.header), guard.value) {
		return nil
	}
	return &HTTPError{StatusCode: http.StatusProxyAuthRequired, Cause: errors.New("proxy API key authentication failed")}
}

type AnyGuard struct {
	guards []Guard
}

func NewAnyGuard(guards ...Guard) *AnyGuard {
	return &AnyGuard{guards: append([]Guard(nil), guards...)}
}

func (guard *AnyGuard) Check(ctx context.Context, request *http.Request) error {
	var lastErr error
	for _, candidate := range guard.guards {
		if candidate == nil {
			continue
		}
		if err := candidate.Check(ctx, request); err == nil {
			return nil
		} else {
			lastErr = err
		}
	}
	if lastErr == nil {
		lastErr = errors.New("no authentication method is configured")
	}
	return lastErr
}

type ClientIPResolver struct {
	trusted policy.CIDRMatcher
}

func NewClientIPResolver(trustedCIDRs []string) (*ClientIPResolver, error) {
	trusted, err := policy.NewCIDRMatcher(trustedCIDRs)
	if err != nil {
		return nil, fmt.Errorf("create client IP resolver: %w", err)
	}
	return &ClientIPResolver{trusted: trusted}, nil
}

func (resolver *ClientIPResolver) Resolve(request *http.Request) (netip.Addr, error) {
	if resolver == nil || request == nil {
		return netip.Addr{}, errors.New("resolve client IP: resolver and request are required")
	}
	peer, err := parseRemoteAddress(request.RemoteAddr)
	if err != nil {
		return netip.Addr{}, err
	}
	if !resolver.trusted.Match(peer) {
		return peer, nil
	}

	chain, present, err := parseForwardedChain(request.Header.Values("Forwarded"))
	if err != nil {
		return netip.Addr{}, err
	}
	if !present {
		chain, err = parseXForwardedFor(request.Header.Values("X-Forwarded-For"))
		if err != nil {
			return netip.Addr{}, err
		}
	}
	if len(chain) == 0 {
		return peer, nil
	}
	for index := len(chain) - 1; index >= 0; index-- {
		if !resolver.trusted.Match(chain[index]) {
			return chain[index], nil
		}
	}
	if len(chain) > 0 {
		return chain[0], nil
	}
	return peer, nil
}

func parseXForwardedFor(fields []string) ([]netip.Addr, error) {
	chain := make([]netip.Addr, 0, len(fields)+1)
	for _, field := range fields {
		for value := range strings.SplitSeq(field, ",") {
			address, err := netip.ParseAddr(strings.TrimSpace(value))
			if err != nil {
				return nil, fmt.Errorf("resolve client IP: invalid X-Forwarded-For address %q", value)
			}
			chain = append(chain, address.Unmap())
		}
	}
	return chain, nil
}

func parseForwardedChain(fields []string) ([]netip.Addr, bool, error) {
	if len(fields) == 0 {
		return nil, false, nil
	}
	chain := make([]netip.Addr, 0, len(fields)+1)
	for _, field := range fields {
		for element := range strings.SplitSeq(field, ",") {
			found := false
			for parameter := range strings.SplitSeq(element, ";") {
				name, value, ok := strings.Cut(strings.TrimSpace(parameter), "=")
				if !ok || !strings.EqualFold(name, "for") {
					continue
				}
				address, err := parseForwardedIdentifier(value)
				if err != nil {
					return nil, true, err
				}
				chain = append(chain, address)
				found = true
				break
			}
			if !found {
				return nil, true, errors.New("resolve client IP: Forwarded element is missing for parameter")
			}
		}
	}
	return chain, true, nil
}

func parseForwardedIdentifier(raw string) (netip.Addr, error) {
	value := strings.TrimSpace(raw)
	if strings.HasPrefix(value, `"`) {
		unquoted, err := strconv.Unquote(value)
		if err != nil {
			return netip.Addr{}, fmt.Errorf("resolve client IP: invalid quoted Forwarded identifier")
		}
		value = unquoted
	}
	if strings.EqualFold(value, "unknown") || strings.HasPrefix(value, "_") {
		return netip.Addr{}, fmt.Errorf("resolve client IP: non-IP Forwarded identifier")
	}
	if strings.HasPrefix(value, "[") {
		closing := strings.IndexByte(value, ']')
		if closing < 0 {
			return netip.Addr{}, fmt.Errorf("resolve client IP: invalid Forwarded IPv6 identifier")
		}
		value = value[1:closing]
	} else if host, _, err := net.SplitHostPort(value); err == nil {
		value = host
	}
	address, err := netip.ParseAddr(value)
	if err != nil {
		return netip.Addr{}, fmt.Errorf("resolve client IP: invalid Forwarded address")
	}
	return address.Unmap(), nil
}

type AccessGuard struct {
	resolver     *ClientIPResolver
	allow        policy.CIDRMatcher
	unrestricted bool
}

func NewAccessGuard(resolver *ClientIPResolver, allowCIDRs []string) (*AccessGuard, error) {
	if resolver == nil {
		return nil, errors.New("create access guard: client IP resolver is required")
	}
	allow, err := policy.NewCIDRMatcher(allowCIDRs)
	if err != nil {
		return nil, fmt.Errorf("create access guard: %w", err)
	}
	return &AccessGuard{resolver: resolver, allow: allow, unrestricted: len(allowCIDRs) == 0}, nil
}

func (guard *AccessGuard) Check(_ context.Context, request *http.Request) error {
	address, err := guard.resolver.Resolve(request)
	if err != nil {
		return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
	}
	if guard.unrestricted || guard.allow.Match(address) {
		return nil
	}
	return &HTTPError{StatusCode: http.StatusForbidden, Cause: fmt.Errorf("client address %s is not allowed", address)}
}

type Admitter interface {
	Admit(context.Context, string) error
}

type AdmissionGuard struct {
	resolver *ClientIPResolver
	admitter Admitter
}

func NewAdmissionGuard(resolver *ClientIPResolver, admitter Admitter) *AdmissionGuard {
	return &AdmissionGuard{resolver: resolver, admitter: admitter}
}

func (guard *AdmissionGuard) Check(ctx context.Context, request *http.Request) error {
	if guard == nil || guard.resolver == nil || guard.admitter == nil {
		return &HTTPError{StatusCode: http.StatusInternalServerError, Cause: errors.New("gateway admission is not configured")}
	}
	address, err := guard.resolver.Resolve(request)
	if err != nil {
		return &HTTPError{StatusCode: http.StatusBadRequest, Cause: err}
	}
	if err := guard.admitter.Admit(ctx, address.String()); err != nil {
		return &HTTPError{StatusCode: http.StatusTooManyRequests, Cause: err}
	}
	return nil
}

func parseBasicCredentials(value string) (string, string, bool) {
	scheme, encoded, ok := strings.Cut(strings.TrimSpace(value), " ")
	if !ok || !strings.EqualFold(scheme, "Basic") {
		return "", "", false
	}
	decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(encoded))
	if err != nil {
		return "", "", false
	}
	username, password, ok := strings.Cut(string(decoded), ":")
	return username, password, ok
}

func constantTimeEqual(actual, expected string) bool {
	return subtle.ConstantTimeCompare([]byte(actual), []byte(expected)) == 1
}

func parseRemoteAddress(remote string) (netip.Addr, error) {
	host, _, err := net.SplitHostPort(strings.TrimSpace(remote))
	if err != nil {
		host = strings.TrimSpace(remote)
	}
	address, err := netip.ParseAddr(host)
	if err != nil {
		return netip.Addr{}, fmt.Errorf("resolve client IP: invalid remote address %q", remote)
	}
	return address.Unmap(), nil
}