360 lines
12 KiB
Go
360 lines
12 KiB
Go
package controlplane
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"time"
|
|
|
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
"proxy-pool/internal/controlplane/snapshotwire"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
routingDomain "proxy-pool/internal/domain/routing"
|
|
"proxy-pool/internal/domain/workerruntime"
|
|
"proxy-pool/internal/gateway/snapshot"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/protobuf/types/known/emptypb"
|
|
"google.golang.org/protobuf/types/known/timestamppb"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidSnapshotWatcher = errors.New("invalid gateway snapshot watcher")
|
|
ErrSnapshotChecksum = errors.New("worker snapshot checksum mismatch")
|
|
ErrSnapshotDelta = errors.New("worker snapshot delta requires a full snapshot")
|
|
)
|
|
|
|
type SnapshotStream interface {
|
|
Recv() (*controlplanev1.SnapshotEnvelope, error)
|
|
}
|
|
|
|
type SnapshotRPCClient interface {
|
|
Watch(context.Context, *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error)
|
|
Acknowledge(context.Context, *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error)
|
|
}
|
|
|
|
type SnapshotWatcherOptions struct {
|
|
ClusterID string
|
|
WorkerID string
|
|
}
|
|
|
|
type SnapshotWatcher struct {
|
|
client SnapshotRPCClient
|
|
store *snapshot.Store
|
|
options SnapshotWatcherOptions
|
|
}
|
|
|
|
func NewSnapshotWatcher(client SnapshotRPCClient, store *snapshot.Store, options SnapshotWatcherOptions) (*SnapshotWatcher, error) {
|
|
if client == nil || store == nil || !workerruntime.ValidIdentifier(options.ClusterID) || !workerruntime.ValidIdentifier(options.WorkerID) {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
return &SnapshotWatcher{client: client, store: store, options: options}, nil
|
|
}
|
|
|
|
// Watch applies verified full snapshots in stream order. Delta support remains
|
|
// intentionally fail-closed until the Gateway can validate a complete routing view.
|
|
func (watcher *SnapshotWatcher) Watch(ctx context.Context, sessionID string) error {
|
|
if watcher == nil || ctx == nil || !workerruntime.ValidIdentifier(sessionID) {
|
|
return ErrInvalidSnapshotWatcher
|
|
}
|
|
request := &controlplanev1.WatchSnapshotsRequest{WorkerId: watcher.options.WorkerID, SessionId: sessionID}
|
|
if current := watcher.store.Current(); current != nil {
|
|
request.LastAppliedVersion = current.Version
|
|
}
|
|
stream, err := watcher.client.Watch(ctx, request)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for {
|
|
envelope, err := stream.Recv()
|
|
if errors.Is(err, io.EOF) {
|
|
return nil
|
|
}
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := watcher.applyAndAcknowledge(ctx, sessionID, envelope); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
|
|
func (watcher *SnapshotWatcher) applyAndAcknowledge(ctx context.Context, sessionID string, envelope *controlplanev1.SnapshotEnvelope) error {
|
|
if envelope == nil || envelope.GetFull() == nil {
|
|
return ErrSnapshotDelta
|
|
}
|
|
full := envelope.GetFull()
|
|
applyErr := watcher.applyFull(full)
|
|
if len(full.GetChecksum()) != sha256.Size {
|
|
return errors.Join(applyErr, ErrSnapshotChecksum)
|
|
}
|
|
acknowledgement := &controlplanev1.AcknowledgeSnapshotRequest{
|
|
WorkerId: watcher.options.WorkerID, SessionId: sessionID, Version: full.GetVersion(),
|
|
OwnershipEpoch: full.GetOwnershipEpoch(), Checksum: append([]byte(nil), full.GetChecksum()...),
|
|
Applied: applyErr == nil,
|
|
}
|
|
if applyErr != nil {
|
|
acknowledgement.ErrorCode = "snapshot_apply_failed"
|
|
}
|
|
_, acknowledgeErr := watcher.client.Acknowledge(ctx, acknowledgement)
|
|
if applyErr != nil || acknowledgeErr != nil {
|
|
return errors.Join(applyErr, acknowledgeErr)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (watcher *SnapshotWatcher) applyFull(full *controlplanev1.WorkerSnapshot) error {
|
|
if full == nil || full.GetVersion() == 0 || full.GetOwnershipEpoch() == 0 {
|
|
return ErrInvalidSnapshotWatcher
|
|
}
|
|
validUntil, err := requiredFutureTimestamp(full.GetValidUntil(), time.Now().UTC())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(full.GetChecksum()) != sha256.Size || string(checksum[:]) != string(full.GetChecksum()) {
|
|
return ErrSnapshotChecksum
|
|
}
|
|
proxies, err := wireProxies(full.GetProxies())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
routes, err := wireRouting(full.GetRouting())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
envelope := snapshot.Envelope{
|
|
ClusterID: watcher.options.ClusterID, WorkerID: watcher.options.WorkerID,
|
|
Epoch: full.GetOwnershipEpoch(), Version: full.GetVersion(), Full: true, ValidUntil: validUntil,
|
|
Proxies: proxies, Routing: routes,
|
|
}
|
|
envelope.Checksum = snapshot.ChecksumWithRouting(proxies, routes)
|
|
return watcher.store.Apply(envelope)
|
|
}
|
|
|
|
func wireRouting(source []*controlplanev1.RoutingRule) ([]routingDomain.Rule, error) {
|
|
result := make([]routingDomain.Rule, 0, len(source))
|
|
names := make(map[string]struct{}, len(source))
|
|
for _, rule := range source {
|
|
if rule == nil || !workerruntime.ValidIdentifier(rule.GetName()) {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
if _, duplicate := names[rule.GetName()]; duplicate {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
names[rule.GetName()] = struct{}{}
|
|
if !rule.GetEnabled() {
|
|
continue
|
|
}
|
|
upstreams, err := wireRoutingUpstreams(rule.GetUpstreams())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
strategy, err := wireRoutingStrategy(rule.GetStrategy(), upstreams)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
action, err := wireUnavailableAction(rule.GetOnUnavailable())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
result = append(result, routingDomain.Rule{
|
|
Name: rule.GetName(),
|
|
Match: routingDomain.Match{
|
|
HostRegex: rule.GetHostRegex(), Methods: append([]string(nil), rule.GetMethods()...),
|
|
PathRegex: rule.GetPathRegex(), Headers: cloneRoutingHeaders(rule.GetHeaders()),
|
|
},
|
|
Upstreams: upstreams, Action: routingDomain.ActionProxy, Strategy: strategy, OnUnavailable: action,
|
|
})
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func wireRoutingUpstreams(source []string) ([]string, error) {
|
|
if len(source) == 0 {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
result := make([]string, len(source))
|
|
seen := make(map[string]struct{}, len(source))
|
|
for index, upstream := range source {
|
|
if !workerruntime.ValidIdentifier(upstream) {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
if _, duplicate := seen[upstream]; duplicate {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
seen[upstream] = struct{}{}
|
|
result[index] = upstream
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func wireRoutingStrategy(source *controlplanev1.RoutingStrategy, upstreams []string) (routingDomain.Strategy, error) {
|
|
if source == nil {
|
|
return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
strategy := routingDomain.Strategy{CurrentUpstream: source.GetCurrentUpstream()}
|
|
switch source.GetType() {
|
|
case controlplanev1.StrategyType_STRATEGY_TYPE_SEQUENTIAL:
|
|
strategy.Type = routingDomain.StrategySequential
|
|
if !containsUpstream(upstreams, strategy.CurrentUpstream) {
|
|
return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
case controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM:
|
|
strategy.Type = routingDomain.StrategyRandom
|
|
case controlplanev1.StrategyType_STRATEGY_TYPE_ROUND_ROBIN:
|
|
strategy.Type = routingDomain.StrategyRoundRobin
|
|
case controlplanev1.StrategyType_STRATEGY_TYPE_WEIGHTED:
|
|
strategy.Type = routingDomain.StrategyWeighted
|
|
strategy.Weights = make(map[string]uint32, len(upstreams))
|
|
for _, upstream := range upstreams {
|
|
weight, exists := source.GetWeights()[upstream]
|
|
if !exists || weight == 0 {
|
|
return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
strategy.Weights[upstream] = weight
|
|
}
|
|
for upstream := range source.GetWeights() {
|
|
if !containsUpstream(upstreams, upstream) {
|
|
return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
}
|
|
case controlplanev1.StrategyType_STRATEGY_TYPE_LEAST_CONNECTIONS:
|
|
strategy.Type = routingDomain.StrategyLeastConnections
|
|
default:
|
|
return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
return strategy, nil
|
|
}
|
|
|
|
func wireUnavailableAction(action controlplanev1.UnavailableAction) (routingDomain.OnUnavailableAction, error) {
|
|
switch action {
|
|
case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT:
|
|
return routingDomain.OnUnavailableReject, nil
|
|
case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_WAIT:
|
|
return routingDomain.OnUnavailableWait, nil
|
|
case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_DIRECT:
|
|
return routingDomain.OnUnavailableDirect, nil
|
|
default:
|
|
return "", fmt.Errorf("%w: unsupported unavailable action", ErrInvalidSnapshotWatcher)
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
func wireProxies(source []*controlplanev1.OwnedProxy) ([]proxyDomain.Proxy, error) {
|
|
proxies := make([]proxyDomain.Proxy, len(source))
|
|
for index, item := range source {
|
|
proxy, err := wireProxy(item)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
proxies[index] = proxy
|
|
}
|
|
return proxies, nil
|
|
}
|
|
|
|
func wireProxy(source *controlplanev1.OwnedProxy) (proxyDomain.Proxy, error) {
|
|
if source == nil || !workerruntime.ValidIdentifier(source.GetId()) || !workerruntime.ValidIdentifier(source.GetUpstream()) ||
|
|
source.GetHost() == "" || source.GetPort() == 0 || source.GetPort() > math.MaxUint16 || source.GetMaxConcurrency() == 0 {
|
|
return proxyDomain.Proxy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
scheme, ok := wireScheme(source.GetProtocol())
|
|
if !ok {
|
|
return proxyDomain.Proxy{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
expiresAt, err := wireTimestamp(source.GetExpiresAt())
|
|
if err != nil {
|
|
return proxyDomain.Proxy{}, err
|
|
}
|
|
usableUntil, err := wireTimestamp(source.GetUsableUntil())
|
|
if err != nil {
|
|
return proxyDomain.Proxy{}, err
|
|
}
|
|
tags := make(map[string]string, len(source.GetTags()))
|
|
for key, value := range source.GetTags() {
|
|
tags[key] = value
|
|
}
|
|
return proxyDomain.Proxy{
|
|
ID: source.GetId(), Scheme: scheme, Host: source.GetHost(), Port: uint16(source.GetPort()),
|
|
Username: source.GetUsername(), CredentialVersion: source.GetCredentialVersion(), SecretRef: source.GetSecretRef(),
|
|
SourceUpstream: source.GetUpstream(), ExpiresAt: expiresAt, UsableUntil: usableUntil,
|
|
MaxConcurrency: int64(source.GetMaxConcurrency()), State: proxyDomain.StateAvailable, Tags: tags,
|
|
}, nil
|
|
}
|
|
|
|
func wireScheme(protocol controlplanev1.ProxyProtocol) (proxyDomain.Scheme, bool) {
|
|
switch protocol {
|
|
case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP:
|
|
return proxyDomain.SchemeHTTP, true
|
|
case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS:
|
|
return proxyDomain.SchemeHTTPS, true
|
|
case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5:
|
|
return proxyDomain.SchemeSOCKS5, true
|
|
default:
|
|
return "", false
|
|
}
|
|
}
|
|
|
|
func wireTimestamp(value *timestamppb.Timestamp) (*time.Time, error) {
|
|
if value == nil {
|
|
return nil, nil
|
|
}
|
|
if err := value.CheckValid(); err != nil {
|
|
return nil, ErrInvalidSnapshotWatcher
|
|
}
|
|
converted := value.AsTime().UTC()
|
|
return &converted, nil
|
|
}
|
|
|
|
func requiredFutureTimestamp(value *timestamppb.Timestamp, now time.Time) (time.Time, error) {
|
|
if value == nil || value.CheckValid() != nil || now.IsZero() {
|
|
return time.Time{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
converted := value.AsTime().UTC()
|
|
if !converted.After(now) {
|
|
return time.Time{}, ErrInvalidSnapshotWatcher
|
|
}
|
|
return converted, nil
|
|
}
|
|
|
|
type generatedSnapshotRPCClient struct {
|
|
client controlplanev1.WorkerControlPlaneClient
|
|
}
|
|
|
|
func NewGeneratedSnapshotRPCClient(client controlplanev1.WorkerControlPlaneClient) SnapshotRPCClient {
|
|
return generatedSnapshotRPCClient{client: client}
|
|
}
|
|
|
|
func (client generatedSnapshotRPCClient) Watch(ctx context.Context, request *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) {
|
|
return client.client.WatchSnapshots(ctx, request, grpc.WaitForReady(true))
|
|
}
|
|
|
|
func (client generatedSnapshotRPCClient) Acknowledge(ctx context.Context, request *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
|
|
return client.client.AcknowledgeSnapshot(ctx, request, grpc.WaitForReady(true))
|
|
}
|