proxy-pool/internal/controller/worker/grpc_handler.go

package worker

import (
	"context"
	"crypto/sha256"
	"errors"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/domain/workerruntime"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/emptypb"
)

type IdentityAuthorizer interface {
	Authorize(context.Context, string) error
}

type GRPCHandler struct {
	controlplanev1.UnimplementedWorkerControlPlaneServer
	service   Service
	identity  IdentityAuthorizer
	snapshots SnapshotSource
}

func NewGRPCHandler(service Service, identity IdentityAuthorizer, snapshots ...SnapshotSource) *GRPCHandler {
	handler := &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) {
	if request == nil || handler == nil || handler.service == nil || handler.identity == nil {
		return nil, grpcError(ErrInvalidCommand)
	}
	if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
		return nil, err
	}
	registration, err := handler.service.Register(ctx, RegisterCommand{
		WorkerID: request.GetWorkerId(), InstanceID: request.GetInstanceId(), Zone: request.GetZone(),
		ProtocolVersion: request.GetSupportedProtocolVersion(), Labels: cloneLabels(request.GetLabels()),
	})
	if err != nil {
		return nil, grpcError(err)
	}
	return &controlplanev1.RegisterWorkerResponse{
		WorkerId: registration.WorkerID, SessionId: registration.SessionID, OwnershipEpoch: registration.OwnershipEpoch,
		HeartbeatInterval: durationpb.New(registration.HeartbeatInterval), MaxStaleAge: durationpb.New(registration.MaxStaleAge),
	}, 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)
	}
	if err := handler.service.ValidateSession(stream.Context(), request.GetWorkerId(), request.GetSessionId()); err != nil {
		return grpcError(err)
	}
	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(), request.GetSessionId(), 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, sessionID 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, sessionID, 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) {
	if request == nil || handler == nil || handler.service == nil || handler.identity == nil {
		return nil, grpcError(ErrInvalidCommand)
	}
	if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
		return nil, err
	}
	err := handler.service.Acknowledge(ctx, SnapshotAcknowledgement{
		WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), Version: request.GetVersion(),
		OwnershipEpoch: request.GetOwnershipEpoch(), Checksum: append([]byte(nil), request.GetChecksum()...),
		Applied: request.GetApplied(), ErrorCode: request.GetErrorCode(), ErrorMessage: request.GetErrorMessage(),
	})
	if err != nil {
		return nil, grpcError(err)
	}
	return &emptypb.Empty{}, nil
}

func (handler *GRPCHandler) ReportRuntime(ctx context.Context, request *controlplanev1.ReportRuntimeRequest) (*controlplanev1.ReportRuntimeResponse, error) {
	if request == nil || handler == nil || handler.service == nil || handler.identity == nil || request.GetObservedAt() == nil || request.GetObservedAt().CheckValid() != nil {
		return nil, grpcError(ErrInvalidCommand)
	}
	if err := handler.authorize(ctx, request.GetWorkerId()); err != nil {
		return nil, err
	}
	counters := make([]workerruntime.Counter, len(request.GetCounters()))
	for index, counter := range request.GetCounters() {
		if counter == nil {
			return nil, grpcError(ErrInvalidCommand)
		}
		counters[index] = workerruntime.Counter{
			ProxyID: counter.GetProxyId(), Active: int64(counter.GetActive()), Reserved: int64(counter.GetReserved()), Draining: counter.GetDraining(),
		}
	}
	decision, err := handler.service.ReportRuntime(ctx, workerruntime.Report{
		WorkerID: request.GetWorkerId(), SessionID: request.GetSessionId(), Sequence: request.GetReportSequence(),
		SnapshotVersion: request.GetSnapshotVersion(), OwnershipEpoch: request.GetOwnershipEpoch(),
		ObservedAt: request.GetObservedAt().AsTime(), Counters: counters,
	})
	if err != nil {
		return nil, grpcError(err)
	}
	return &controlplanev1.ReportRuntimeResponse{
		AcceptedOwnershipEpoch: decision.AcceptedOwnershipEpoch, RequireFullSnapshot: decision.RequireFullSnapshot,
	}, nil
}

func (handler *GRPCHandler) authorize(ctx context.Context, workerID string) error {
	if err := handler.identity.Authorize(ctx, workerID); err != nil {
		return status.Error(codes.PermissionDenied, "worker identity is not authorized")
	}
	return nil
}

func grpcError(err error) error {
	switch {
	case errors.Is(err, context.Canceled):
		return status.Error(codes.Canceled, "worker control request canceled")
	case errors.Is(err, context.DeadlineExceeded):
		return status.Error(codes.DeadlineExceeded, "worker control request deadline exceeded")
	case errors.Is(err, ErrInvalidCommand), errors.Is(err, workerruntime.ErrInvalidReport),
		errors.Is(err, workerruntime.ErrInvalidAcknowledgement), errors.Is(err, workerruntime.ErrInvalidSnapshotReference):
		return status.Error(codes.InvalidArgument, "invalid worker control request")
	case errors.Is(err, ErrProtocolVersion):
		return status.Error(codes.FailedPrecondition, "unsupported worker protocol version")
	case errors.Is(err, workerruntime.ErrStaleSession):
		return status.Error(codes.FailedPrecondition, "worker session is stale")
	case errors.Is(err, workerruntime.ErrSnapshotMismatch):
		return status.Error(codes.FailedPrecondition, "worker snapshot does not match issued snapshot")
	case errors.Is(err, workerruntime.ErrStaleAcknowledgement):
		return status.Error(codes.Aborted, "worker snapshot acknowledgement is stale")
	case errors.Is(err, workerruntime.ErrStaleReport):
		return status.Error(codes.Aborted, "worker runtime sequence is stale")
	case errors.Is(err, workerruntime.ErrConflictingReport):
		return status.Error(codes.AlreadyExists, "worker runtime sequence conflicts")
	default:
		return status.Error(codes.Unavailable, "worker control plane unavailable")
	}
}

func cloneLabels(labels map[string]string) map[string]string {
	result := make(map[string]string, len(labels))
	for key, value := range labels {
		result[key] = value
	}
	return result
}