From f28009c0989a4fd8ea25b7e4792de1a902656702 Mon Sep 17 00:00:00 2001 From: youfak Date: Fri, 31 Jul 2026 13:13:19 +0800 Subject: [PATCH] feat: add gateway snapshot watcher client --- internal/gateway/controlplane/watcher.go | 223 ++++++++++++++++++ internal/gateway/controlplane/watcher_test.go | 97 ++++++++ 2 files changed, 320 insertions(+) create mode 100644 internal/gateway/controlplane/watcher.go create mode 100644 internal/gateway/controlplane/watcher_test.go diff --git a/internal/gateway/controlplane/watcher.go b/internal/gateway/controlplane/watcher.go new file mode 100644 index 0000000..79c9888 --- /dev/null +++ b/internal/gateway/controlplane/watcher.go @@ -0,0 +1,223 @@ +package controlplane + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "io" + "math" + "time" + + controlplanev1 "proxy-pool/gen/controlplane/v1" + proxyDomain "proxy-pool/internal/domain/proxy" + "proxy-pool/internal/domain/workerruntime" + "proxy-pool/internal/gateway/snapshot" + + "google.golang.org/grpc" + "google.golang.org/protobuf/proto" + "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 + } + checksum, err := workerSnapshotChecksum(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 + } + envelope := snapshot.Envelope{ + ClusterID: watcher.options.ClusterID, WorkerID: watcher.options.WorkerID, + Epoch: full.GetOwnershipEpoch(), Version: full.GetVersion(), Full: true, Proxies: proxies, + } + envelope.Checksum = snapshot.Checksum(proxies) + return watcher.store.Apply(envelope) +} + +func workerSnapshotChecksum(full *controlplanev1.WorkerSnapshot) ([sha256.Size]byte, error) { + if full == nil { + return [sha256.Size]byte{}, ErrInvalidSnapshotWatcher + } + copy := proto.Clone(full).(*controlplanev1.WorkerSnapshot) + copy.Checksum = nil + encoded, err := proto.MarshalOptions{Deterministic: true}.Marshal(copy) + if err != nil { + return [sha256.Size]byte{}, fmt.Errorf("marshal worker snapshot: %w", err) + } + return sha256.Sum256(encoded), nil +} + +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 +} + +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)) +} diff --git a/internal/gateway/controlplane/watcher_test.go b/internal/gateway/controlplane/watcher_test.go new file mode 100644 index 0000000..96fe697 --- /dev/null +++ b/internal/gateway/controlplane/watcher_test.go @@ -0,0 +1,97 @@ +package controlplane + +import ( + "context" + "crypto/sha256" + "io" + "testing" + "time" + + controlplanev1 "proxy-pool/gen/controlplane/v1" + "proxy-pool/internal/gateway/snapshot" + + "google.golang.org/protobuf/types/known/emptypb" + "google.golang.org/protobuf/types/known/timestamppb" +) + +func TestSnapshotWatcherAppliesVerifiedFullSnapshotAndAcknowledges(t *testing.T) { + store := snapshot.NewStore("cluster-a", "worker-a") + full := &controlplanev1.WorkerSnapshot{ + Version: 1, OwnershipEpoch: 7, GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)), + Proxies: []*controlplanev1.OwnedProxy{{ + Id: "proxy-a", Upstream: "upstream-a", Protocol: controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP, + Host: "192.0.2.10", Port: 8080, MaxConcurrency: 3, ExpiresAt: timestamppb.New(time.Now().Add(time.Minute)), + }}, + } + setSnapshotChecksum(t, full) + client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}} + watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"}) + if err != nil { + t.Fatalf("NewSnapshotWatcher(): %v", err) + } + if err := watcher.Watch(context.Background(), "session-a"); err != nil { + t.Fatalf("Watch(): %v", err) + } + view := store.Current() + if view == nil || view.Version != 1 || view.Epoch != 7 || len(view.Entries) != 1 || view.Entries[0].Proxy.ID != "proxy-a" { + t.Fatalf("snapshot view = %+v", view) + } + if client.watch.GetSessionId() != "session-a" || client.ack.GetVersion() != 1 || !client.ack.GetApplied() || string(client.ack.GetChecksum()) != string(full.GetChecksum()) { + t.Fatalf("watch=%+v ack=%+v", client.watch, client.ack) + } +} + +func TestSnapshotWatcherRejectsChecksumAndAcknowledgesFailure(t *testing.T) { + store := snapshot.NewStore("cluster-a", "worker-a") + full := &controlplanev1.WorkerSnapshot{Version: 1, OwnershipEpoch: 7, Checksum: make([]byte, sha256.Size)} + client := &snapshotClientStub{stream: &snapshotStreamStub{values: []*controlplanev1.SnapshotEnvelope{{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: full}}}}} + watcher, err := NewSnapshotWatcher(client, store, SnapshotWatcherOptions{ClusterID: "cluster-a", WorkerID: "worker-a"}) + if err != nil { + t.Fatalf("NewSnapshotWatcher(): %v", err) + } + if err := watcher.Watch(context.Background(), "session-a"); err == nil { + t.Fatal("Watch() error = nil, want checksum rejection") + } + if client.ack == nil || client.ack.GetApplied() || client.ack.GetVersion() != 1 { + t.Fatalf("negative acknowledgement = %+v", client.ack) + } +} + +type snapshotClientStub struct { + stream SnapshotStream + watch *controlplanev1.WatchSnapshotsRequest + ack *controlplanev1.AcknowledgeSnapshotRequest +} + +func (client *snapshotClientStub) Watch(_ context.Context, request *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) { + client.watch = request + return client.stream, nil +} + +func (client *snapshotClientStub) Acknowledge(_ context.Context, acknowledgement *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) { + client.ack = acknowledgement + return &emptypb.Empty{}, nil +} + +type snapshotStreamStub struct { + values []*controlplanev1.SnapshotEnvelope + index int +} + +func (stream *snapshotStreamStub) Recv() (*controlplanev1.SnapshotEnvelope, error) { + if stream.index >= len(stream.values) { + return nil, io.EOF + } + value := stream.values[stream.index] + stream.index++ + return value, nil +} + +func setSnapshotChecksum(t *testing.T, full *controlplanev1.WorkerSnapshot) { + t.Helper() + checksum, err := workerSnapshotChecksum(full) + if err != nil { + t.Fatalf("workerSnapshotChecksum(): %v", err) + } + full.Checksum = checksum[:] +}