proxy-pool/internal/controller/worker/routing_snapshot_source.go

package worker

import (
	"context"
	"errors"
	"fmt"
	"math"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/config"
	"proxy-pool/internal/domain/adminstate"

	"google.golang.org/protobuf/types/known/durationpb"
)

var ErrRoutingUnavailable = errors.New("worker snapshot routing is unavailable")

// RoutingSource reads the complete, ordered Gateway routing payload for a
// WorkerSnapshot. It never performs data-plane reads.
type RoutingSource interface {
	Read(context.Context) ([]*controlplanev1.RoutingRule, error)
}

type RoutingConfigurationReader interface {
	Snapshot() (*config.Config, uint64)
}

type RoutingStateReader interface {
	Snapshot(context.Context) (adminstate.Snapshot, error)
}

// GatewayRoutingSource combines immutable configuration semantics with the
// current Admin routing and Upstream state. Admin state is optional only for
// deployments that do not enable the Admin control surface.
type GatewayRoutingSource struct {
	configuration RoutingConfigurationReader
	state         RoutingStateReader
}

func NewGatewayRoutingSource(configuration RoutingConfigurationReader, states ...RoutingStateReader) (*GatewayRoutingSource, error) {
	if configuration == nil || len(states) > 1 || len(states) == 1 && states[0] == nil {
		return nil, ErrRoutingUnavailable
	}
	source := &GatewayRoutingSource{configuration: configuration}
	if len(states) == 1 {
		source.state = states[0]
	}
	return source, nil
}

func (source *GatewayRoutingSource) Read(ctx context.Context) ([]*controlplanev1.RoutingRule, error) {
	if source == nil || source.configuration == nil || ctx == nil {
		return nil, ErrRoutingUnavailable
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	configuration, revision := source.configuration.Snapshot()
	if configuration == nil {
		return nil, ErrRoutingUnavailable
	}
	var state *adminstate.Snapshot
	if source.state != nil {
		current, err := source.state.Snapshot(ctx)
		if err != nil {
			return nil, errors.Join(ErrRoutingUnavailable, err)
		}
		if current.Config == nil || current.Config.Revision != revision {
			return nil, ErrRoutingUnavailable
		}
		state = &current
	}
	rules, err := buildGatewayRouting(configuration, state)
	if err != nil {
		return nil, errors.Join(ErrRoutingUnavailable, err)
	}
	return rules, nil
}

func buildGatewayRouting(configuration *config.Config, state *adminstate.Snapshot) ([]*controlplanev1.RoutingRule, error) {
	if configuration == nil {
		return nil, errors.New("configuration is nil")
	}
	upstreamEnabled, err := enabledGatewayUpstreams(configuration, state)
	if err != nil {
		return nil, err
	}
	routingStates, err := gatewayRoutingStates(configuration, state)
	if err != nil {
		return nil, err
	}
	rules := make([]*controlplanev1.RoutingRule, 0, len(configuration.Routing))
	for _, configured := range configuration.Routing {
		if configured.Purpose != "gateway" {
			continue
		}
		state, hasState := routingStates[configured.Name]
		if !hasState && routingStates != nil {
			return nil, fmt.Errorf("routing %q is missing from admin state", configured.Name)
		}
		routeAction, err := wireRoutingAction(configured.Action)
		if err != nil {
			return nil, fmt.Errorf("routing %q action: %w", configured.Name, err)
		}
		enabled := configured.Enabled && (routingStates == nil || state.Enabled)
		if routeAction != controlplanev1.RoutingAction_ROUTING_ACTION_PROXY {
			rules = append(rules, &controlplanev1.RoutingRule{
				Name: configured.Name, Enabled: enabled, HostRegex: configured.Match.HostRegex,
				Methods: append([]string(nil), configured.Match.Methods...), PathRegex: configured.Match.PathRegex,
				Headers: cloneRoutingHeaders(configured.Match.Headers), Action: routeAction,
			})
			continue
		}
		upstreams := enabledRoutingUpstreams(configured.Upstreams, upstreamEnabled)
		current := configuredCurrentUpstream(configured, state, hasState)
		strategy, err := wireRoutingStrategy(configured.Strategy, current, upstreams)
		if err != nil {
			return nil, fmt.Errorf("routing %q strategy: %w", configured.Name, err)
		}
		enabled = enabled && len(upstreams) > 0
		if strategy.GetType() == controlplanev1.StrategyType_STRATEGY_TYPE_SEQUENTIAL && !containsUpstream(upstreams, current) {
			enabled = false
		}
		unavailableAction, err := wireUnavailableAction(configured.OnUnavailable.Action)
		if err != nil {
			return nil, fmt.Errorf("routing %q unavailable action: %w", configured.Name, err)
		}
		var waitTimeout *durationpb.Duration
		if unavailableAction == controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_WAIT {
			waitTimeout = durationpb.New(configured.OnUnavailable.WaitTimeout.Value())
		}
		rules = append(rules, &controlplanev1.RoutingRule{
			Name: configured.Name, Enabled: enabled, HostRegex: configured.Match.HostRegex,
			Methods: append([]string(nil), configured.Match.Methods...), PathRegex: configured.Match.PathRegex,
			Headers: cloneRoutingHeaders(configured.Match.Headers), Upstreams: upstreams, Strategy: strategy,
			OnUnavailable: unavailableAction, WaitTimeout: waitTimeout, Action: controlplanev1.RoutingAction_ROUTING_ACTION_PROXY,
		})
	}
	return rules, nil
}

func enabledGatewayUpstreams(configuration *config.Config, state *adminstate.Snapshot) (map[string]bool, error) {
	adminStates := make(map[string]adminstate.UpstreamState)
	if state != nil {
		for _, upstream := range state.Upstreams {
			if _, duplicate := adminStates[upstream.Name]; duplicate {
				return nil, fmt.Errorf("duplicate upstream state %q", upstream.Name)
			}
			adminStates[upstream.Name] = upstream
		}
	}
	result := make(map[string]bool, len(configuration.Upstreams))
	for name, upstream := range configuration.Upstreams {
		result[name] = upstream.Enabled
		if state == nil {
			continue
		}
		admin, exists := adminStates[name]
		if !exists {
			return nil, fmt.Errorf("upstream %q is missing from admin state", name)
		}
		result[name] = result[name] && admin.Enabled
	}
	return result, nil
}

func gatewayRoutingStates(configuration *config.Config, state *adminstate.Snapshot) (map[string]adminstate.RoutingState, error) {
	if state == nil {
		return nil, nil
	}
	result := make(map[string]adminstate.RoutingState, len(state.Routings))
	for _, routing := range state.Routings {
		if _, duplicate := result[routing.Name]; duplicate {
			return nil, fmt.Errorf("duplicate routing state %q", routing.Name)
		}
		result[routing.Name] = routing
	}
	for _, routing := range configuration.Routing {
		if routing.Purpose == "gateway" {
			if _, exists := result[routing.Name]; !exists {
				return nil, fmt.Errorf("routing %q is missing from admin state", routing.Name)
			}
		}
	}
	return result, nil
}

func enabledRoutingUpstreams(source []string, enabled map[string]bool) []string {
	result := make([]string, 0, len(source))
	for _, upstream := range source {
		if enabled[upstream] {
			result = append(result, upstream)
		}
	}
	return result
}

func configuredCurrentUpstream(routing config.Routing, state adminstate.RoutingState, hasState bool) string {
	if hasState {
		return state.CurrentUpstream
	}
	if len(routing.Upstreams) == 0 {
		return ""
	}
	return routing.Upstreams[0]
}

func wireRoutingStrategy(configured config.Strategy, current string, upstreams []string) (*controlplanev1.RoutingStrategy, error) {
	strategy := &controlplanev1.RoutingStrategy{CurrentUpstream: current}
	switch configured.Type {
	case "sequential":
		strategy.Type = controlplanev1.StrategyType_STRATEGY_TYPE_SEQUENTIAL
	case "random":
		strategy.Type = controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM
	case "roundRobin":
		strategy.Type = controlplanev1.StrategyType_STRATEGY_TYPE_ROUND_ROBIN
	case "weighted":
		strategy.Type = controlplanev1.StrategyType_STRATEGY_TYPE_WEIGHTED
		strategy.Weights = make(map[string]uint32, len(upstreams))
		for _, upstream := range upstreams {
			weight, exists := configured.Weights[upstream]
			if !exists || weight <= 0 || weight > math.MaxUint32 {
				return nil, fmt.Errorf("invalid weight for upstream %q", upstream)
			}
			strategy.Weights[upstream] = uint32(weight)
		}
	case "leastConnections":
		strategy.Type = controlplanev1.StrategyType_STRATEGY_TYPE_LEAST_CONNECTIONS
	default:
		return nil, fmt.Errorf("unsupported strategy type %q", configured.Type)
	}
	return strategy, nil
}

func wireUnavailableAction(action string) (controlplanev1.UnavailableAction, error) {
	switch action {
	case "reject":
		return controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT, nil
	case "wait":
		return controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_WAIT, nil
	case "direct":
		return controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_DIRECT, nil
	default:
		return controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_UNSPECIFIED, fmt.Errorf("unsupported action %q", action)
	}
}

func wireRoutingAction(action string) (controlplanev1.RoutingAction, error) {
	switch action {
	case "", "proxy":
		return controlplanev1.RoutingAction_ROUTING_ACTION_PROXY, nil
	case "direct":
		return controlplanev1.RoutingAction_ROUTING_ACTION_DIRECT, nil
	case "reject":
		return controlplanev1.RoutingAction_ROUTING_ACTION_REJECT, nil
	default:
		return controlplanev1.RoutingAction_ROUTING_ACTION_UNSPECIFIED, fmt.Errorf("unsupported action %q", action)
	}
}

func cloneRoutingHeaders(source map[string]string) map[string]string {
	if source == nil {
		return nil
	}
	result := make(map[string]string, len(source))
	for name, value := range source {
		result[name] = value
	}
	return result
}

func containsUpstream(upstreams []string, target string) bool {
	for _, upstream := range upstreams {
		if upstream == target {
			return true
		}
	}
	return false
}

var _ RoutingSource = (*GatewayRoutingSource)(nil)