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

package worker

import (
	"context"
	"errors"
	"math"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/controlplane/snapshotwire"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	proxyDomain "proxy-pool/internal/domain/proxy"
	"proxy-pool/internal/domain/workerruntime"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var ErrSnapshotCredentialsUnavailable = errors.New("worker snapshot credential material is unavailable")

// OwnedSnapshotSource builds a complete Worker view from the Redis ownership
// index. It only accepts a complete bounded result from the reader.
type OwnedSnapshotSource struct {
	epochs     OwnershipEpochReader
	reader     ownershipDomain.SnapshotReader
	validFor   time.Duration
	maxProxies int
	maxBytes   int
	now        func() time.Time
	routing    RoutingSource
}

func NewOwnedSnapshotSource(
	epochs OwnershipEpochReader,
	reader ownershipDomain.SnapshotReader,
	validFor time.Duration,
	maxProxies int,
	maxBytes int,
	now func() time.Time,
	routing ...RoutingSource,
) (*OwnedSnapshotSource, error) {
	if epochs == nil || reader == nil || validFor <= 0 || maxProxies <= 0 || maxBytes <= 0 || now == nil || len(routing) > 1 ||
		len(routing) == 1 && routing[0] == nil {
		return nil, ErrSnapshotsUnavailable
	}
	source := &OwnedSnapshotSource{
		epochs: epochs, reader: reader, validFor: validFor, maxProxies: maxProxies, maxBytes: maxBytes, now: now,
	}
	if len(routing) == 1 {
		source.routing = routing[0]
	}
	return source, nil
}

func (source *OwnedSnapshotSource) Watch(ctx context.Context, request SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) {
	if source == nil || ctx == nil || !workerruntime.ValidIdentifier(request.WorkerID) ||
		!workerruntime.ValidIdentifier(request.SessionID) || request.LastAppliedVersion == math.MaxUint64 {
		return nil, ErrSnapshotsUnavailable
	}
	epoch, err := source.epochs.CurrentOwnershipEpoch(ctx)
	if err != nil {
		return nil, err
	}
	if epoch == 0 {
		return nil, ErrSnapshotsUnavailable
	}
	now := source.now().UTC()
	if now.IsZero() {
		return nil, ErrSnapshotsUnavailable
	}
	proxies, err := source.reader.ReadWorkerSnapshot(ctx, request.WorkerID, source.maxProxies)
	if err != nil {
		return nil, err
	}
	full := &controlplanev1.WorkerSnapshot{
		Version: request.LastAppliedVersion + 1, OwnershipEpoch: epoch,
		GeneratedAt: timestamppb.New(now), ValidUntil: timestamppb.New(now.Add(source.validFor)),
		Proxies: make([]*controlplanev1.OwnedProxy, 0, len(proxies)),
	}
	for _, item := range proxies {
		owned, validUntil, err := wireOwnedProxy(item, now)
		if err != nil {
			return nil, err
		}
		if validUntil.Before(full.ValidUntil.AsTime()) {
			full.ValidUntil = timestamppb.New(validUntil)
		}
		full.Proxies = append(full.Proxies, owned)
	}
	if source.routing != nil {
		routing, err := source.routing.Read(ctx)
		if err != nil {
			return nil, err
		}
		full.Routing = cloneSnapshotRouting(routing)
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil {
		return nil, err
	}
	full.Checksum = append([]byte(nil), checksum[:]...)
	if proto.Size(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}) > source.maxBytes {
		return nil, ErrSnapshotsUnavailable
	}
	updates := make(chan *controlplanev1.WorkerSnapshot, 1)
	updates <- full
	return updates, nil
}

func cloneSnapshotRouting(source []*controlplanev1.RoutingRule) []*controlplanev1.RoutingRule {
	result := make([]*controlplanev1.RoutingRule, len(source))
	for index, rule := range source {
		if rule != nil {
			result[index] = proto.Clone(rule).(*controlplanev1.RoutingRule)
		}
	}
	return result
}

func wireOwnedProxy(item ownershipDomain.SnapshotProxy, now time.Time) (*controlplanev1.OwnedProxy, time.Time, error) {
	proxy := item.Proxy
	if item.OwnershipEpoch == 0 || item.LeaseExpiresAt.IsZero() || proxy.CredentialVersion != "" || proxy.SecretRef != "" {
		if proxy.CredentialVersion != "" || proxy.SecretRef != "" {
			return nil, time.Time{}, ErrSnapshotCredentialsUnavailable
		}
		return nil, time.Time{}, ErrSnapshotsUnavailable
	}
	protocol, ok := wireProtocol(proxy.Scheme)
	if !ok || !workerruntime.ValidIdentifier(proxy.ID) || !workerruntime.ValidIdentifier(proxy.SourceUpstream) ||
		proxy.Host == "" || proxy.Port == 0 || proxy.MaxConcurrency <= 0 || proxy.MaxConcurrency > math.MaxUint32 ||
		proxy.ExpiresAt == nil || proxy.UsableUntil == nil {
		return nil, time.Time{}, ErrSnapshotsUnavailable
	}
	validUntil := proxy.UsableUntil.UTC()
	if item.LeaseExpiresAt.Before(validUntil) {
		validUntil = item.LeaseExpiresAt.UTC()
	}
	if !validUntil.After(now) || !proxy.ExpiresAt.After(now) {
		return nil, time.Time{}, ErrSnapshotsUnavailable
	}
	tags := make(map[string]string, len(proxy.Tags))
	for key, value := range proxy.Tags {
		tags[key] = value
	}
	return &controlplanev1.OwnedProxy{
		Id: proxy.ID, Upstream: proxy.SourceUpstream, Protocol: protocol, Host: proxy.Host, Port: uint32(proxy.Port),
		Username: proxy.Username, ExpiresAt: timestamppb.New(proxy.ExpiresAt.UTC()),
		MaxConcurrency: uint32(proxy.MaxConcurrency), Tags: tags, OwnershipEpoch: item.OwnershipEpoch,
		UsableUntil: timestamppb.New(validUntil),
	}, validUntil, nil
}

func wireProtocol(scheme proxyDomain.Scheme) (controlplanev1.ProxyProtocol, bool) {
	switch scheme {
	case proxyDomain.SchemeHTTP:
		return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP, true
	case proxyDomain.SchemeHTTPS:
		return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS, true
	case proxyDomain.SchemeSOCKS5:
		return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5, true
	default:
		return controlplanev1.ProxyProtocol_PROXY_PROTOCOL_UNSPECIFIED, false
	}
}

var _ SnapshotSource = (*OwnedSnapshotSource)(nil)