gitee.com/h79/goutils@v1.22.10/rpc/auth.go (about)

     1  package rpc
     2  
     3  import (
     4  	"context"
     5  	"google.golang.org/grpc"
     6  )
     7  
     8  type AuthFunc func(ctx context.Context, req interface{}, fullMethodName string) (context.Context, error)
     9  
    10  // ServiceAuthOverride allows a given gRPC service implementation to override the global `AuthFunc`.
    11  //
    12  // If a service implements the AuthFuncOverride method, it takes precedence over the `AuthFunc` method,
    13  // and will be called instead of AuthFunc for all method invocations within that service.
    14  type ServiceAuthOverride interface {
    15  	AuthFuncOverride(ctx context.Context, req interface{}, fullMethodName string) (context.Context, error)
    16  }
    17  
    18  // UnaryServerInterceptor returns a new unary server interceptors that performs per-request auth.
    19  func UnaryServerInterceptor(authFunc AuthFunc) grpc.UnaryServerInterceptor {
    20  	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    21  		var newCtx context.Context
    22  		var err error
    23  		if overrideSrv, ok := info.Server.(ServiceAuthOverride); ok {
    24  			newCtx, err = overrideSrv.AuthFuncOverride(ctx, req, info.FullMethod)
    25  		} else {
    26  			newCtx, err = authFunc(ctx, req, info.FullMethod)
    27  		}
    28  		if err != nil {
    29  			return nil, err
    30  		}
    31  		return handler(newCtx, req)
    32  	}
    33  }