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)