241 lines
8.5 KiB
Go
241 lines
8.5 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"math"
|
|
"sort"
|
|
"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"
|
|
platformCredentials "proxy-pool/internal/platform/credentials"
|
|
|
|
"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
|
|
credentials CredentialReader
|
|
}
|
|
|
|
type CredentialReader interface {
|
|
Resolve(context.Context, platformCredentials.Reference) (platformCredentials.Value, error)
|
|
}
|
|
|
|
func NewOwnedSnapshotSource(
|
|
epochs OwnershipEpochReader,
|
|
reader ownershipDomain.SnapshotReader,
|
|
validFor time.Duration,
|
|
maxProxies int,
|
|
maxBytes int,
|
|
now func() time.Time,
|
|
routing ...RoutingSource,
|
|
) (*OwnedSnapshotSource, error) {
|
|
return newOwnedSnapshotSource(epochs, reader, validFor, maxProxies, maxBytes, now, nil, routing...)
|
|
}
|
|
|
|
func NewOwnedSnapshotSourceWithCredentials(
|
|
epochs OwnershipEpochReader,
|
|
reader ownershipDomain.SnapshotReader,
|
|
validFor time.Duration,
|
|
maxProxies int,
|
|
maxBytes int,
|
|
now func() time.Time,
|
|
credentials CredentialReader,
|
|
routing ...RoutingSource,
|
|
) (*OwnedSnapshotSource, error) {
|
|
if credentials == nil {
|
|
return nil, ErrSnapshotsUnavailable
|
|
}
|
|
return newOwnedSnapshotSource(epochs, reader, validFor, maxProxies, maxBytes, now, credentials, routing...)
|
|
}
|
|
|
|
func newOwnedSnapshotSource(
|
|
epochs OwnershipEpochReader,
|
|
reader ownershipDomain.SnapshotReader,
|
|
validFor time.Duration,
|
|
maxProxies int,
|
|
maxBytes int,
|
|
now func() time.Time,
|
|
credentials CredentialReader,
|
|
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,
|
|
credentials: credentials,
|
|
}
|
|
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)),
|
|
}
|
|
materials := make(map[platformCredentials.Reference]platformCredentials.Value)
|
|
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 owned.GetSecretRef() != "" {
|
|
reference := platformCredentials.Reference{SecretRef: owned.GetSecretRef(), CredentialVersion: owned.GetCredentialVersion()}
|
|
if source.credentials == nil {
|
|
return nil, ErrSnapshotCredentialsUnavailable
|
|
}
|
|
material, resolveErr := source.credentials.Resolve(ctx, reference)
|
|
if resolveErr != nil || (material.Username == "" && material.Password == "") ||
|
|
(owned.GetUsername() != "" && material.Username != "" && material.Username != owned.GetUsername()) {
|
|
return nil, ErrSnapshotCredentialsUnavailable
|
|
}
|
|
if material.Username == "" {
|
|
material.Username = owned.GetUsername()
|
|
}
|
|
if existing, exists := materials[reference]; exists && existing != material {
|
|
return nil, ErrSnapshotCredentialsUnavailable
|
|
}
|
|
materials[reference] = material
|
|
}
|
|
}
|
|
full.Credentials = wireSnapshotCredentials(materials)
|
|
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 wireSnapshotCredentials(materials map[platformCredentials.Reference]platformCredentials.Value) []*controlplanev1.SnapshotCredential {
|
|
result := make([]*controlplanev1.SnapshotCredential, 0, len(materials))
|
|
for reference, material := range materials {
|
|
result = append(result, &controlplanev1.SnapshotCredential{
|
|
SecretRef: reference.SecretRef, CredentialVersion: reference.CredentialVersion,
|
|
Username: material.Username, Password: material.Password,
|
|
})
|
|
}
|
|
sort.Slice(result, func(left, right int) bool {
|
|
if result[left].GetSecretRef() == result[right].GetSecretRef() {
|
|
return result[left].GetCredentialVersion() < result[right].GetCredentialVersion()
|
|
}
|
|
return result[left].GetSecretRef() < result[right].GetSecretRef()
|
|
})
|
|
return result
|
|
}
|
|
|
|
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, CredentialVersion: proxy.CredentialVersion, SecretRef: proxy.SecretRef, 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)
|