Compare commits
No commits in common. "8b734b85f35c054a09207e67220f548e0f12b082" and "de2c9ce9b6a11e8325de385617cbc8dc2172928b" have entirely different histories.
8b734b85f3
...
de2c9ce9b6
@ -2,13 +2,11 @@ package worker
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"crypto/sha256"
|
|
||||||
"errors"
|
"errors"
|
||||||
|
|
||||||
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
||||||
"proxy-pool/internal/domain/workerruntime"
|
"proxy-pool/internal/domain/workerruntime"
|
||||||
|
|
||||||
"google.golang.org/grpc"
|
|
||||||
"google.golang.org/grpc/codes"
|
"google.golang.org/grpc/codes"
|
||||||
"google.golang.org/grpc/status"
|
"google.golang.org/grpc/status"
|
||||||
"google.golang.org/protobuf/types/known/durationpb"
|
"google.golang.org/protobuf/types/known/durationpb"
|
||||||
@ -23,15 +21,10 @@ type GRPCHandler struct {
|
|||||||
controlplanev1.UnimplementedWorkerControlPlaneServer
|
controlplanev1.UnimplementedWorkerControlPlaneServer
|
||||||
service Service
|
service Service
|
||||||
identity IdentityAuthorizer
|
identity IdentityAuthorizer
|
||||||
snapshots SnapshotSource
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewGRPCHandler(service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) *GRPCHandler {
|
func NewGRPCHandler(service Service, identity IdentityAuthorizer) *GRPCHandler {
|
||||||
handler := &GRPCHandler{service: service, identity: identity}
|
return &GRPCHandler{service: service, identity: identity}
|
||||||
if len(snapshots) == 1 {
|
|
||||||
handler.snapshots = snapshots[0]
|
|
||||||
}
|
|
||||||
return handler
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (handler *GRPCHandler) RegisterWorker(ctx context.Context, request *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
|
func (handler *GRPCHandler) RegisterWorker(ctx context.Context, request *controlplanev1.RegisterWorkerRequest) (*controlplanev1.RegisterWorkerResponse, error) {
|
||||||
@ -54,58 +47,6 @@ func (handler *GRPCHandler) RegisterWorker(ctx context.Context, request *control
|
|||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (handler *GRPCHandler) WatchSnapshots(request *controlplanev1.WatchSnapshotsRequest, stream grpc.ServerStreamingServer[controlplanev1.SnapshotEnvelope]) error {
|
|
||||||
if request == nil || stream == nil || handler == nil || handler.service == nil || handler.identity == nil {
|
|
||||||
return grpcError(ErrInvalidCommand)
|
|
||||||
}
|
|
||||||
if handler.snapshots == nil {
|
|
||||||
return status.Error(codes.Unimplemented, "worker snapshots are unavailable")
|
|
||||||
}
|
|
||||||
if err := handler.authorize(stream.Context(), request.GetWorkerId()); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
if !workerruntime.ValidIdentifier(request.GetWorkerId()) || !workerruntime.ValidIdentifier(request.GetSessionId()) ||
|
|
||||||
(len(request.GetLastChecksum()) != 0 && len(request.GetLastChecksum()) != sha256.Size) {
|
|
||||||
return grpcError(ErrInvalidCommand)
|
|
||||||
}
|
|
||||||
updates, err := handler.snapshots.Watch(stream.Context(), SnapshotWatchRequest{
|
|
||||||
WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), LastAppliedVersion: request.GetLastAppliedVersion(),
|
|
||||||
LastChecksum: append([]byte(nil), request.GetLastChecksum()...),
|
|
||||||
})
|
|
||||||
if err != nil {
|
|
||||||
return grpcError(err)
|
|
||||||
}
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case <-stream.Context().Done():
|
|
||||||
return stream.Context().Err()
|
|
||||||
case snapshot, ok := <-updates:
|
|
||||||
if !ok {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
if err := handler.issueSnapshot(stream.Context(), request.GetWorkerId(), snapshot); err != nil {
|
|
||||||
return grpcError(err)
|
|
||||||
}
|
|
||||||
if err := stream.Send(&controlplanev1.SnapshotEnvelope{Payload: &controlplanev1.SnapshotEnvelope_Full{Full: snapshot}}); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (handler *GRPCHandler) issueSnapshot(ctx context.Context, workerID string, snapshot *controlplanev1.WorkerSnapshot) error {
|
|
||||||
if snapshot == nil || snapshot.GetVersion() == 0 || snapshot.GetOwnershipEpoch() == 0 || len(snapshot.GetChecksum()) != sha256.Size ||
|
|
||||||
snapshot.GetGeneratedAt() == nil || snapshot.GetGeneratedAt().CheckValid() != nil ||
|
|
||||||
snapshot.GetValidUntil() == nil || snapshot.GetValidUntil().CheckValid() != nil {
|
|
||||||
return ErrInvalidCommand
|
|
||||||
}
|
|
||||||
var checksum [sha256.Size]byte
|
|
||||||
copy(checksum[:], snapshot.GetChecksum())
|
|
||||||
return handler.service.IssueSnapshot(ctx, workerruntime.SnapshotReference{
|
|
||||||
WorkerID: workerID, Version: snapshot.GetVersion(), OwnershipEpoch: snapshot.GetOwnershipEpoch(), Checksum: checksum,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (handler *GRPCHandler) AcknowledgeSnapshot(ctx context.Context, request *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
|
func (handler *GRPCHandler) AcknowledgeSnapshot(ctx context.Context, request *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
|
||||||
if request == nil || handler == nil || handler.service == nil || handler.identity == nil {
|
if request == nil || handler == nil || handler.service == nil || handler.identity == nil {
|
||||||
return nil, grpcError(ErrInvalidCommand)
|
return nil, grpcError(ErrInvalidCommand)
|
||||||
|
|||||||
@ -64,29 +64,6 @@ func TestGRPCHandlerMapsErrorsAndLeavesStreamsUnimplemented(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestGRPCHandlerStreamsIssuedFullSnapshots(t *testing.T) {
|
|
||||||
checksum := make([]byte, 32)
|
|
||||||
checksum[0] = 1
|
|
||||||
full := &controlplanev1.WorkerSnapshot{
|
|
||||||
Version: 3, OwnershipEpoch: 9, Checksum: checksum,
|
|
||||||
GeneratedAt: timestamppb.New(time.Now()), ValidUntil: timestamppb.New(time.Now().Add(time.Minute)),
|
|
||||||
}
|
|
||||||
service := &grpcServiceStub{}
|
|
||||||
client, cleanup := grpcWorkerClient(t, service, allowIdentity{}, snapshotSourceStub{snapshots: []*controlplanev1.WorkerSnapshot{full}})
|
|
||||||
defer cleanup()
|
|
||||||
stream, err := client.WatchSnapshots(context.Background(), &controlplanev1.WatchSnapshotsRequest{WorkerId: "worker-a", SessionId: "session-a"})
|
|
||||||
if err != nil {
|
|
||||||
t.Fatalf("WatchSnapshots(): %v", err)
|
|
||||||
}
|
|
||||||
received, err := stream.Recv()
|
|
||||||
if err != nil || received.GetFull().GetVersion() != 3 {
|
|
||||||
t.Fatalf("Recv() = %+v, %v", received, err)
|
|
||||||
}
|
|
||||||
if service.issued.WorkerID != "worker-a" || service.issued.Version != 3 || service.issued.Checksum[0] != 1 {
|
|
||||||
t.Fatalf("issued snapshot = %+v", service.issued)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
type grpcServiceStub struct {
|
type grpcServiceStub struct {
|
||||||
registration Registration
|
registration Registration
|
||||||
registerErr error
|
registerErr error
|
||||||
@ -94,16 +71,11 @@ type grpcServiceStub struct {
|
|||||||
acknowledgeErr error
|
acknowledgeErr error
|
||||||
report workerruntime.Report
|
report workerruntime.Report
|
||||||
reportErr error
|
reportErr error
|
||||||
issued workerruntime.SnapshotReference
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) {
|
func (stub *grpcServiceStub) Register(context.Context, RegisterCommand) (Registration, error) {
|
||||||
return stub.registration, stub.registerErr
|
return stub.registration, stub.registerErr
|
||||||
}
|
}
|
||||||
func (stub *grpcServiceStub) IssueSnapshot(_ context.Context, reference workerruntime.SnapshotReference) error {
|
|
||||||
stub.issued = reference
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error {
|
func (stub *grpcServiceStub) Acknowledge(_ context.Context, acknowledgement SnapshotAcknowledgement) error {
|
||||||
stub.acknowledgement = acknowledgement
|
stub.acknowledgement = acknowledgement
|
||||||
return stub.acknowledgeErr
|
return stub.acknowledgeErr
|
||||||
@ -117,24 +89,11 @@ type allowIdentity struct{}
|
|||||||
|
|
||||||
func (allowIdentity) Authorize(context.Context, string) error { return nil }
|
func (allowIdentity) Authorize(context.Context, string) error { return nil }
|
||||||
|
|
||||||
type snapshotSourceStub struct {
|
func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer) (controlplanev1.WorkerControlPlaneClient, func()) {
|
||||||
snapshots []*controlplanev1.WorkerSnapshot
|
|
||||||
}
|
|
||||||
|
|
||||||
func (source snapshotSourceStub) Watch(_ context.Context, _ SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error) {
|
|
||||||
updates := make(chan *controlplanev1.WorkerSnapshot, len(source.snapshots))
|
|
||||||
for _, snapshot := range source.snapshots {
|
|
||||||
updates <- snapshot
|
|
||||||
}
|
|
||||||
close(updates)
|
|
||||||
return updates, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func grpcWorkerClient(t *testing.T, service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) (controlplanev1.WorkerControlPlaneClient, func()) {
|
|
||||||
t.Helper()
|
t.Helper()
|
||||||
listener := bufconn.Listen(1 << 20)
|
listener := bufconn.Listen(1 << 20)
|
||||||
server := grpc.NewServer()
|
server := grpc.NewServer()
|
||||||
controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity, snapshots...))
|
controlplanev1.RegisterWorkerControlPlaneServer(server, NewGRPCHandler(service, identity))
|
||||||
go func() { _ = server.Serve(listener) }()
|
go func() { _ = server.Serve(listener) }()
|
||||||
connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithInsecure())
|
connection, err := grpc.NewClient("passthrough:///bufnet", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { return listener.Dial() }), grpc.WithInsecure())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
|||||||
@ -23,7 +23,6 @@ var ErrInvalidServer = errors.New("invalid worker control server configuration")
|
|||||||
|
|
||||||
type ServerOptions struct {
|
type ServerOptions struct {
|
||||||
ShutdownTimeout time.Duration
|
ShutdownTimeout time.Duration
|
||||||
Snapshots SnapshotSource
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func DefaultServerOptions() ServerOptions {
|
func DefaultServerOptions() ServerOptions {
|
||||||
@ -61,7 +60,7 @@ func NewServer(controlPlane config.ControlPlane, service Service, options Server
|
|||||||
}),
|
}),
|
||||||
)
|
)
|
||||||
grpcServer := grpc.NewServer(serverOptions...)
|
grpcServer := grpc.NewServer(serverOptions...)
|
||||||
controlplanev1.RegisterWorkerControlPlaneServer(grpcServer, NewGRPCHandler(service, identity, options.Snapshots))
|
controlplanev1.RegisterWorkerControlPlaneServer(grpcServer, NewGRPCHandler(service, identity))
|
||||||
return &Server{listen: controlPlane.Listen, grpcServer: grpcServer, shutdownTimeout: options.ShutdownTimeout}, nil
|
return &Server{listen: controlPlane.Listen, grpcServer: grpcServer, shutdownTimeout: options.ShutdownTimeout}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -60,28 +60,10 @@ type Options struct {
|
|||||||
|
|
||||||
type Service interface {
|
type Service interface {
|
||||||
Register(context.Context, RegisterCommand) (Registration, error)
|
Register(context.Context, RegisterCommand) (Registration, error)
|
||||||
IssueSnapshot(context.Context, workerruntime.SnapshotReference) error
|
|
||||||
Acknowledge(context.Context, SnapshotAcknowledgement) error
|
Acknowledge(context.Context, SnapshotAcknowledgement) error
|
||||||
ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error)
|
ReportRuntime(context.Context, workerruntime.Report) (RuntimeDecision, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (service *service) IssueSnapshot(ctx context.Context, reference workerruntime.SnapshotReference) error {
|
|
||||||
if ctx == nil {
|
|
||||||
return ErrInvalidCommand
|
|
||||||
}
|
|
||||||
if err := ctx.Err(); err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
normalized, err := workerruntime.NormalizeSnapshotReference(reference)
|
|
||||||
if err != nil {
|
|
||||||
return errors.Join(ErrInvalidCommand, err)
|
|
||||||
}
|
|
||||||
if err := service.store.RecordIssuedSnapshot(ctx, normalized, service.options.SessionTTL); err != nil {
|
|
||||||
return classifyStoreError(err)
|
|
||||||
}
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
type service struct {
|
type service struct {
|
||||||
store workerruntime.ControlStore
|
store workerruntime.ControlStore
|
||||||
options Options
|
options Options
|
||||||
|
|||||||
@ -35,8 +35,8 @@ func TestServiceRegistersAcknowledgesAndReportsRuntime(t *testing.T) {
|
|||||||
WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch,
|
WorkerID: "worker-a", Version: 7, OwnershipEpoch: registered.OwnershipEpoch,
|
||||||
Checksum: sha256.Sum256([]byte("snapshot-7")),
|
Checksum: sha256.Sum256([]byte("snapshot-7")),
|
||||||
}
|
}
|
||||||
if err := service.IssueSnapshot(context.Background(), reference); err != nil {
|
if err := store.RecordIssuedSnapshot(context.Background(), reference, time.Minute); err != nil {
|
||||||
t.Fatalf("IssueSnapshot(): %v", err)
|
t.Fatalf("RecordIssuedSnapshot(): %v", err)
|
||||||
}
|
}
|
||||||
if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{
|
if err := service.Acknowledge(context.Background(), SnapshotAcknowledgement{
|
||||||
WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7,
|
WorkerID: "worker-a", SessionID: registered.SessionID, Version: 7,
|
||||||
|
|||||||
@ -1,21 +0,0 @@
|
|||||||
package worker
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
|
|
||||||
controlplanev1 "proxy-pool/gen/controlplane/v1"
|
|
||||||
)
|
|
||||||
|
|
||||||
var ErrSnapshotsUnavailable = errors.New("worker snapshots are unavailable")
|
|
||||||
|
|
||||||
type SnapshotWatchRequest struct {
|
|
||||||
WorkerID string
|
|
||||||
SessionID string
|
|
||||||
LastAppliedVersion uint64
|
|
||||||
LastChecksum []byte
|
|
||||||
}
|
|
||||||
|
|
||||||
type SnapshotSource interface {
|
|
||||||
Watch(context.Context, SnapshotWatchRequest) (<-chan *controlplanev1.WorkerSnapshot, error)
|
|
||||||
}
|
|
||||||
@ -1,223 +0,0 @@
|
|||||||
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))
|
|
||||||
}
|
|
||||||
@ -1,97 +0,0 @@
|
|||||||
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[:]
|
|
||||||
}
|
|
||||||
Loading…
Reference in New Issue
Block a user