utils/interceptor/v1/grpc/interceptor.go

package grpc

import (
	"context"
	"gitea.drugeyes.vip/pharnexbase/utils/enum"
	"github.com/go-kratos/kratos/v2/errors"
	"github.com/go-kratos/kratos/v2/log"
	"google.golang.org/grpc"
)

var ErrInternalServer = errors.InternalServer("InternalServerError", "服务错误")

type Interceptor struct {
	env    enum.Env         // 环境
	reason map[string]int32 // 业务错误枚举
}

func NewInterceptor(env enum.Env, reason map[string]int32) *Interceptor {
	return &Interceptor{
		env:    env,
		reason: reason,
	}
}

func (i *Interceptor) UnaryClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		h := func(ctx context.Context, req interface{}) (interface{}, error) {
			return reply, invoker(ctx, method, req, reply, cc, opts...)
		}
		_, err := h(ctx, req)
		if err != nil {
			err = errors.FromError(err)
		}
		return err
	}
}

func (i *Interceptor) UnaryServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		resp, err = handler(ctx, req)
		if err == nil {
			return
		}

		if i.env == enum.Env_Production {
			se := errors.FromError(err)
			if _, ok := i.reason[se.Reason]; ok {
				// 业务错误，原样返回
				return
			}
			// 正式服，统一返回“服务错误”
			err = ErrInternalServer
			log.Errorf("InterServerErr:%v", se)
		}
		return
	}
}