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

package health

import (
	"context"
	"errors"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/domain/activitypool"
	healthDomain "proxy-pool/internal/domain/health"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

var ErrInvalidGRPCHandler = errors.New("invalid checker grpc handler")

// CheckerIdentityAuthorizer proves that the calling mTLS identity owns the
// checker ID in the request. It is deliberately separate from worker session
// identity because Checkers do not own gateway snapshots.
type CheckerIdentityAuthorizer interface {
	AuthorizeChecker(context.Context, string) error
}

type GRPCHandlerOptions struct {
	MaxObservationsPerBatch int
}

func DefaultGRPCHandlerOptions() GRPCHandlerOptions {
	return GRPCHandlerOptions{MaxObservationsPerBatch: 1_000}
}

// GRPCHandler exposes fact reporting now. StreamCheckTasks remains inherited
// as unimplemented until the Controller's leased task scheduler is available.
type GRPCHandler struct {
	controlplanev1.UnimplementedCheckerControlPlaneServer
	reducer  *Reducer
	identity CheckerIdentityAuthorizer
	options  GRPCHandlerOptions
}

func NewGRPCHandler(
	reducer *Reducer,
	identity CheckerIdentityAuthorizer,
	options GRPCHandlerOptions,
) (*GRPCHandler, error) {
	if reducer == nil || nilInterface(identity) {
		return nil, ErrInvalidGRPCHandler
	}
	if options.MaxObservationsPerBatch == 0 {
		options = DefaultGRPCHandlerOptions()
	}
	if options.MaxObservationsPerBatch < 0 {
		return nil, ErrInvalidGRPCHandler
	}
	return &GRPCHandler{reducer: reducer, identity: identity, options: options}, nil
}

func (handler *GRPCHandler) ReportObservations(
	ctx context.Context,
	request *controlplanev1.ObservationBatch,
) (*controlplanev1.ReportObservationsResponse, error) {
	if ctx == nil || handler == nil || handler.reducer == nil || nilInterface(handler.identity) || request == nil ||
		request.GetCheckerId() == "" || len(request.GetObservations()) == 0 ||
		len(request.GetObservations()) > handler.options.MaxObservationsPerBatch {
		return nil, healthGRPCError(ErrInvalidGRPCHandler)
	}
	if err := handler.identity.AuthorizeChecker(ctx, request.GetCheckerId()); err != nil {
		return nil, status.Error(codes.PermissionDenied, "checker identity is not authorized")
	}
	response := &controlplanev1.ReportObservationsResponse{}
	for _, item := range request.GetObservations() {
		observation, err := decodeHealthObservation(item)
		if err == nil {
			_, err = handler.reducer.Apply(ctx, observation)
		}
		if err == nil {
			response.Accepted++
			continue
		}
		if rejectedObservationError(err) {
			response.Rejected++
			continue
		}
		return nil, healthGRPCError(err)
	}
	return response, nil
}

func decodeHealthObservation(item *controlplanev1.HealthObservation) (healthDomain.Observation, error) {
	if item == nil || item.GetLatency() == nil || item.GetLatency().CheckValid() != nil ||
		item.GetObservedAt() == nil || item.GetObservedAt().CheckValid() != nil {
		return healthDomain.Observation{}, healthDomain.ErrInvalidObservation
	}
	level, ok := decodeCheckLevel(item.GetLevel())
	if !ok {
		return healthDomain.Observation{}, healthDomain.ErrInvalidObservation
	}
	return healthDomain.Observation{
		TaskID: item.GetTaskId(), ProxyID: item.GetProxyId(), Level: level,
		RoutingName: item.GetRoutingName(), TargetURL: item.GetTargetUrl(), Success: item.GetSuccess(),
		FailureClass: item.GetFailureClass(), Latency: item.GetLatency().AsDuration(),
		ObservedEgressIP: item.GetObservedEgressIp(), ObservedAt: item.GetObservedAt().AsTime(),
	}, nil
}

func decodeCheckLevel(level controlplanev1.CheckLevel) (healthDomain.Level, bool) {
	switch level {
	case controlplanev1.CheckLevel_CHECK_LEVEL_BASIC:
		return healthDomain.LevelBasic, true
	case controlplanev1.CheckLevel_CHECK_LEVEL_EGRESS:
		return healthDomain.LevelEgress, true
	case controlplanev1.CheckLevel_CHECK_LEVEL_TARGET:
		return healthDomain.LevelTarget, true
	default:
		return "", false
	}
}

func rejectedObservationError(err error) bool {
	return errors.Is(err, healthDomain.ErrInvalidObservation) || errors.Is(err, healthDomain.ErrInvalidFailureThreshold) ||
		errors.Is(err, healthDomain.ErrStaleObservation) || errors.Is(err, healthDomain.ErrConflictingObservation) ||
		errors.Is(err, activitypool.ErrInvalidHealthUpdate) || errors.Is(err, activitypool.ErrActivityNotFound) ||
		errors.Is(err, ErrInvalidThreshold)
}

func healthGRPCError(err error) error {
	switch {
	case errors.Is(err, context.Canceled):
		return status.Error(codes.Canceled, "checker control request canceled")
	case errors.Is(err, context.DeadlineExceeded):
		return status.Error(codes.DeadlineExceeded, "checker control request deadline exceeded")
	case errors.Is(err, ErrInvalidGRPCHandler), errors.Is(err, healthDomain.ErrInvalidObservation),
		errors.Is(err, healthDomain.ErrInvalidFailureThreshold), errors.Is(err, activitypool.ErrInvalidHealthUpdate),
		errors.Is(err, ErrInvalidThreshold):
		return status.Error(codes.InvalidArgument, "invalid checker observation batch")
	default:
		return status.Error(codes.Unavailable, "checker control plane unavailable")
	}
}