go-zero(九) 自定义拦截器

go-zero 拦截器

有时我们需要在处理请求的过程中添加一些额外的逻辑，比如身份验证、日志记录、请求限流、性能监控等，这些都可以通过拦截器实现。go zero可以设置多个拦截器

一、 服务端拦截器

服务端拦截器用于处理传入的 RPC 请求，可以在处理请求之前或之后自定义逻辑，主要用于请求参数校验、身份验证等。

AddUnaryInterceptors 是在 gRPC 服务器初始化时调用的方法，用于将一个或多个服务端拦截器注册到 gRPC 服务器上.

下面我们就简单的演示下日志输出，在rpc项目中，修改 main 函数：

func main() {
	flag.Parse()

	var c config.Config
	conf.MustLoad(*configFile, &c)
	ctx := svc.NewServiceContext(c)

	s := zrpc.MustNewServer(c.RpcServerConf, func(grpcServer *grpc.Server) {
		pb.RegisterUserServer(grpcServer, server.NewUserServer(ctx))

		if c.Mode == service.DevMode || c.Mode == service.TestMode {
			reflection.Register(grpcServer)
		}
	})
	defer s.Stop()
	s.AddUnaryInterceptors(TestUnaryInterceptors)  //添加服务端拦截器

	fmt.Printf("Starting rpc server at %s...\n", c.ListenOn)
	s.Start()
}

//定义拦截器 
func TestUnaryInterceptors(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
	fmt.Printf("req :%v\n", req)
	fmt.Printf("info :%v\n", info)
	return handler(ctx, req)
}

运行项目测试：

二、客户端拦截器

客户端拦截器，用于拦截并处理发送到后端服务的请求，主要用于请求日志记录、性能监测等。

WithUnaryClientInterceptor 是一个用于配置 gRPC 客户端的选项。通过它，可以指定一个或多个拦截器，拦截器会在客户端请求 RPC 方法时被调用。

下面我们演示下使用拦截器输出请求方法耗时，修改api项目的中的servicecontext.go文件：

func NewServiceContext(c config.Config) *ServiceContext {
	return &ServiceContext{
		Config:    c,
		UserModel: model.NewUsersModel(sqlx.NewMysql(c.MysqlDB.DbSource)),

		TestMiddleware: middleware.NewMD5PasswordMiddleware().Handle,
		UserRpc:        user.NewUser(zrpc.MustNewClient(c.UserRpcConf, zrpc.WithUnaryClientInterceptor(TestClientInterceptor))), //在RPC调用的时候添加拦截器
	}
}

func TestClientInterceptor(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
	stime := time.Now()
	// 拦截前
	err := invoker(ctx, method, req, reply, cc, opts...)
	if err != nil {
		return err
	}
	// 拦截后
	fmt.Printf("调用 %s 方法 耗时: %v\n", method, time.Since(stime))
	return nil
}