Python gRPC 服务器拦截器错误返回机制详解

import grpc

# 辅助函数：创建一个终止 RPC 的方法处理器
def _unary_unary_rpc_terminator(code, details):
    """
    生成一个一元RPC方法处理器，该处理器在被调用时会终止RPC
    并返回指定的gRPC状态码和详细信息。
    """
    def terminate(ignored_request, context):
        """
        实际执行终止操作的函数。
        """
        context.abort(code, details)

    # 返回一个 grpc.unary_unary_rpc_method_handler 实例，
    # 它的实际处理逻辑就是上面的 terminate 函数。
    return grpc.unary_unary_rpc_method_handler(terminate)


class RequestHeaderValidatorInterceptor(grpc.ServerInterceptor):
    """
    一个gRPC服务器拦截器示例，用于验证请求头是否存在特定值。
    如果验证失败，则终止请求并返回错误。
    """
    def __init__(self, header_name, expected_value, error_code, error_details):
        """
        初始化拦截器。

        Args:
            header_name (str): 需要验证的请求头名称。
            expected_value (str): 请求头期望的值。
            error_code (grpc.StatusCode): 验证失败时返回的gRPC状态码。
            error_details (str): 验证失败时返回的错误详细信息。
        """
        self._header_name = header_name
        self._expected_value = expected_value
        # 预先创建一个终止处理器，以便在验证失败时直接返回。
        self._terminator = _unary_unary_rpc_terminator(error_code, error_details)

    def intercept_service(self, continuation, handler_call_details):
        """
        拦截服务调用。

        Args:
            continuation: 一个函数，如果拦截器允许请求继续，则调用它。
                          它接受 handler_call_details 作为参数并返回 RpcMethodHandler。
            handler_call_details: 包含有关当前RPC调用的详细信息，
                                  如方法名称和调用元数据（请求头）。

        Returns:
            grpc.RpcMethodHandler: 如果请求合法，则返回 continuation 的结果；
                                   否则，返回预设的终止处理器。
        """
        # 检查请求元数据中是否存在指定的请求头及其值
        # handler_call_details.invocation_metadata 是一个元组列表，例如 [('key', 'value')]
        if (self._header_name, self._expected_value) in handler_call_details.invocation_metadata:
            # 如果验证通过，则允许请求继续执行到实际的服务方法
            return continuation(handler_call_details)
        else:
            # 如果验证失败，则返回预先创建的终止处理器
            # 当gRPC框架尝试调用这个处理器时，它将执行 context.abort()
            return self._terminator

# 示例用法：
# 假设我们有一个gRPC服务器实例 server
# server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))

# 创建一个拦截器，要求请求头中必须有 ('authorization', 'valid-token')
# 否则返回 UNAUTHENTICATED 错误
# auth_interceptor = RequestHeaderValidatorInterceptor(
#     "authorization",
#     "valid-token",
#     grpc.StatusCode.UNAUTHENTICATED,
#     "Authentication failed: Missing or invalid authorization token."
# )

# 将拦截器添加到服务器
# server.intercept_unary_unary(auth_interceptor)
# server.add_insecure_port('[::]:50051')
# server.start()
# server.wait_for_termination()