gRPC微服务搭建（学习阶段2:权限校验）

在gRPC的权限校验中，写一个类，然后继承grpc.ServerInterceptor即可使用权限校验，每隔函数在调用前，都会调用这个函数，这个函数可用于权限校验或是一些数据准备等操作。

文章目录

权限校验示例程序
- 服务端代码
- 客户端代码
带Login的权限管理
- 服务端代码
- 客户端代码

权限校验示例程序

服务端代码

py 复制代码

import json
from concurrent import futures
import grpc


# 业务函数：收到客户端请求后，返回问候消息。
def say_hello(request, _context):
    return {"message": f"Hello, {request['name']}!"}


# 拦截器是 gRPC 的固定扩展写法，可以在请求到达业务函数前做日志、鉴权等处理。
# 这里为了保持简单，只打印请求方法，不做真正的权限校验。
class LogInterceptor(grpc.ServerInterceptor):
    def intercept_service(self, next_step, handler_call_details):
        print(f"收到请求：{handler_call_details.method}")
        return next_step(handler_call_details)


def server(env="dev", address="[::]:50051", max_workers=2):
    """启动 gRPC 服务。

    env：运行环境名称，当前只用于打印提示。
    address：监听地址，格式通常是"主机:端口"。
    max_workers：线程池中的最大工作线程数。
    """
    print(f"当前运行环境：{env}")

    # 这是 gRPC 创建服务端的固定写法。
    grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=max_workers), interceptors=[LogInterceptor()])

    # 下面是本示例的 JSON 请求处理器。
    # 标准 gRPC 项目通常会用 .proto 文件自动生成这部分代码。
    handler = grpc.unary_unary_rpc_method_handler(
        say_hello,
        request_deserializer=lambda data: json.loads(data.decode("utf-8")),
        response_serializer=lambda response: json.dumps(response).encode("utf-8"),
    )

    # 注册服务和方法，这是手动注册 RPC 方法的写法。
    grpc_server.add_generic_rpc_handlers((
        grpc.method_handlers_generic_handler("hello.Greeter", {"SayHello": handler}),))

    # 添加监听端口、启动服务、持续等待，是 gRPC 服务启动的固定流程。
    grpc_server.add_insecure_port(address)
    grpc_server.start()
    print(f"gRPC 服务已启动，监听地址：{address}")
    grpc_server.wait_for_termination()


# 程序从这里开始执行，调用 server() 启动服务。
if __name__ == "__main__":
    server()

客户端代码

py 复制代码

import json
import grpc

# 连接地址必须和 server.py 中的端口一致。
with grpc.insecure_channel("127.0.0.1:50051") as channel:
    # 指定要调用的服务方法：/服务名/方法名。
    say_hello = channel.unary_unary(
        "/hello.Greeter/SayHello",
        # 发送请求前：Python 字典转换为 JSON 字节。
        request_serializer=lambda data: json.dumps(data).encode("utf-8"),
        # 收到响应后：JSON 字节转换为 Python 字典。
        response_deserializer=lambda data: json.loads(data.decode("utf-8")),
    )

    # 发送名称，并获取服务端返回的问候消息。
    response = say_hello({"name": "World"})
    print(response["message"])

带Login的权限管理

使用带login的权限管理，核心在于AuthInterceptor这个类中：

当method == "/auth.Auth/Login"时，直接执行校验
当执行其他的method的时候，直接替换调用check_token这个函数校验token，如果可以，就通过handler.unary_unary(request, context)继续执行，如果不行就通过context.abort(grpc.StatusCode.UNAUTHENTICATED, "请先登录")直接返回，不进行后续的执行

服务端代码

py 复制代码

import json
from concurrent import futures
import grpc

# 服务端保存登录成功后的 Token。
TOKEN = None


# 登录接口不需要 Token，登录成功后生成并保存 Token。
def login(request, _context):
    global TOKEN
    if request["username"] != "admin" or request["password"] != "123456":
        return {"success": False, "message": "用户名或密码错误"}
    TOKEN = "token-123"
    return {"success": True, "token": TOKEN, "message": "登录成功"}


# 只有通过 Token 校验后，才会执行这个业务函数。
def say_hello(request, _context):
    return {"message": f"Hello, {request['name']}!"}


# 将普通 Python 函数注册成 gRPC 方法，并处理 JSON 转换。
def json_handler(function):
    return grpc.unary_unary_rpc_method_handler(
        function,
        request_deserializer=lambda data: json.loads(data.decode("utf-8")),
        response_serializer=lambda data: json.dumps(data).encode("utf-8"),
    )


# 在业务函数执行前检查 Token。
class AuthInterceptor(grpc.ServerInterceptor):
    def intercept_service(self, continuation, handler_call_details):
        method = handler_call_details.method
        handler = continuation(handler_call_details)
        if handler is None:
            return None
        # 登录接口不能检查 Token，否则客户端无法登录。
        if method == "/auth.Auth/Login":
            return handler

        # 其他接口都必须携带正确的 Token。
        def check_token(request, context):
            if request.get("token") != TOKEN:
                context.abort(grpc.StatusCode.UNAUTHENTICATED, "请先登录")
            return handler.unary_unary(request, context)  # Token 正确，继续执行原来的业务函数

        # 用权限检查函数替换原来的业务处理函数。
        return handler._replace(unary_unary=check_token)


def server(address="[::]:50051", max_workers=2):
    """启动 gRPC 服务。"""
    # 创建服务，并安装 AuthInterceptor 权限拦截器
    grpc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=max_workers), interceptors=[AuthInterceptor()])
    # 注册登录接口
    grpc_server.add_generic_rpc_handlers(
        (
            grpc.method_handlers_generic_handler(
                "auth.Auth",
                {"Login": json_handler(login)},
            ),
            # 注册需要登录的问候接口。
            grpc.method_handlers_generic_handler(
                "hello.Greeter",
                {"SayHello": json_handler(say_hello)},
            ),
        )
    )
    grpc_server.add_insecure_port(address)
    grpc_server.start()
    print(f"gRPC 服务已启动，监听地址：{address}")
    grpc_server.wait_for_termination()


if __name__ == "__main__":
    server()

客户端代码

py 复制代码

import json
import grpc

# 连接地址必须和 server.py 中的端口一致。
with grpc.insecure_channel("127.0.0.1:50051") as channel:
    # 第一步：调用登录接口，获取服务端保存的 Token。
    login = channel.unary_unary(
        "/auth.Auth/Login",
        request_serializer=lambda data: json.dumps(data).encode("utf-8"),
        response_deserializer=lambda data: json.loads(data.decode("utf-8")),
    )
    login_result = login({"username": "admin", "password": "123456"})
    print(login_result["message"])

    # 第二步：调用需要登录的 SayHello 接口。
    say_hello = channel.unary_unary(
        "/hello.Greeter/SayHello",
        request_serializer=lambda data: json.dumps(data).encode("utf-8"),
        response_deserializer=lambda data: json.loads(data.decode("utf-8")),
    )

    # 登录成功后，把 Token 和名称一起发送给服务端。
    response = say_hello({"name": "World", "token": login_result["token"]})
    print(response["message"])

    # 第三步：不发送 Token，模拟未登录调用。
    try:
        say_hello({"name": "World"})
    except grpc.RpcError as error:
        print(f"未登录：{error.details()}")