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

package worker

import (
	"context"
	"errors"

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

	"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
}

func NewGRPCHandler(service Service, identity IdentityAuthorizer) *GRPCHandler {
	return &GRPCHandler{service: service, identity: identity}
}

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