proxy-pool/internal/domain/routing/rule.go

package routing

import (
	"fmt"
	"regexp"
	"slices"
	"strings"
)

type Action string

const (
	ActionProxy  Action = "proxy"
	ActionDirect Action = "direct"
	ActionReject Action = "reject"
)

type StrategyType string

const (
	StrategySequential       StrategyType = "sequential"
	StrategyRandom           StrategyType = "random"
	StrategyRoundRobin       StrategyType = "roundRobin"
	StrategyWeighted         StrategyType = "weighted"
	StrategyLeastConnections StrategyType = "leastConnections"
)

type Strategy struct {
	Type            StrategyType
	CurrentUpstream string
	Weights         map[string]uint32
}

type OnUnavailableAction string

const (
	OnUnavailableReject OnUnavailableAction = "reject"
	OnUnavailableWait   OnUnavailableAction = "wait"
	OnUnavailableDirect OnUnavailableAction = "direct"
)

type Match struct {
	HostRegex string
	Methods   []string
	PathRegex string
	Headers   map[string]string
}

type Rule struct {
	Name          string
	Match         Match
	Upstreams     []string
	Action        Action
	Strategy      Strategy
	OnUnavailable OnUnavailableAction
}

type Request struct {
	Host    string
	Method  string
	Path    string
	Headers map[string]string
}

type compiledRule struct {
	rule Rule
	host *regexp.Regexp
	path *regexp.Regexp
}

type RuleSet struct {
	rules []compiledRule
}

func Compile(rules []Rule) (*RuleSet, error) {
	compiled := make([]compiledRule, 0, len(rules))
	for _, rule := range rules {
		if rule.Name == "" {
			return nil, fmt.Errorf("compile routing: rule name is required")
		}
		host, err := regexp.Compile(rule.Match.HostRegex)
		if err != nil {
			return nil, fmt.Errorf("compile routing %q host: %w", rule.Name, err)
		}
		var path *regexp.Regexp
		if rule.Match.PathRegex != "" {
			path, err = regexp.Compile(rule.Match.PathRegex)
			if err != nil {
				return nil, fmt.Errorf("compile routing %q path: %w", rule.Name, err)
			}
		}
		owned := cloneRule(rule)
		if owned.Action == "" && len(owned.Upstreams) > 0 {
			owned.Action = ActionProxy
		}
		compiled = append(compiled, compiledRule{rule: owned, host: host, path: path})
	}
	return &RuleSet{rules: compiled}, nil
}

func (r *RuleSet) Match(request Request) (Rule, bool) {
	if r == nil {
		return Rule{}, false
	}
	host := strings.ToLower(strings.TrimSuffix(request.Host, "."))
	method := strings.ToUpper(request.Method)
	for _, candidate := range r.rules {
		if !candidate.host.MatchString(host) {
			continue
		}
		if len(candidate.rule.Match.Methods) > 0 && !containsFold(candidate.rule.Match.Methods, method) {
			continue
		}
		if candidate.path != nil && !candidate.path.MatchString(request.Path) {
			continue
		}
		if !headersMatch(candidate.rule.Match.Headers, request.Headers) {
			continue
		}
		return cloneRule(candidate.rule), true
	}
	return Rule{}, false
}

func cloneRule(source Rule) Rule {
	cloned := source
	cloned.Match.Methods = append([]string(nil), source.Match.Methods...)
	cloned.Upstreams = append([]string(nil), source.Upstreams...)
	if source.Match.Headers != nil {
		cloned.Match.Headers = make(map[string]string, len(source.Match.Headers))
		for name, value := range source.Match.Headers {
			cloned.Match.Headers[name] = value
		}
	}
	if source.Strategy.Weights != nil {
		cloned.Strategy.Weights = make(map[string]uint32, len(source.Strategy.Weights))
		for upstream, weight := range source.Strategy.Weights {
			cloned.Strategy.Weights[upstream] = weight
		}
	}
	return cloned
}

func containsFold(values []string, target string) bool {
	return slices.ContainsFunc(values, func(value string) bool {
		return strings.EqualFold(value, target)
	})
}

func headersMatch(expected, actual map[string]string) bool {
	for name, value := range expected {
		matched := false
		for actualName, actualValue := range actual {
			if strings.EqualFold(name, actualName) && actualValue == value {
				matched = true
				break
			}
		}
		if !matched {
			return false
		}
	}
	return true
}