feat: add gateway snapshot watcher client

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))
+}

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[:]
+}