go grpc 自定义拦截器

proto文件

syntax = "proto3";
import "google/protobuf/timestamp.proto";
package pb;

option go_package = "./pb";

service ToDoService {
    rpc DoWork (TodoRequest) returns (TodoResponse);
}

enum Week {
    Sunday = 0;
    Monday = 1;
    Tuesday = 2;
    Wednesday = 3;
    Thursday = 4;
    Friday = 5;
    Saturday = 6;
}

message TodoRequest{
    string todo = 1;
    Week week = 2;
    map <string,string> bookMap = 3;   //  BookMap  map[string]string
    google.protobuf.Timestamp doneTime = 4;
}

message TodoResponse {
    bool done = 1;

}

生成grpc文件

protoc user.proto --go_out=./ --go-grpc_out=./ --go-grpc_opt=require_unimplemented_servers=false

 服务端

package main

import (
	"context"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"grpc_enum/pb"
	"net"
	"time"
)

//自定义拦截器
func MyInterceptor  (ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error){
	//从context 获取数据
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		fmt.Println("拦截器 ：metadata nil")
		return nil,status.Error(codes.Unauthenticated, "认证失败[metadata]")
	}

	token , ok := md["token"]
	if !ok {
		fmt.Println("拦截器 ：token nil")
		return nil,status.Error(codes.Unauthenticated, "认证失败[token]")
	}
	fmt.Println("token:",token)

	now := time.Now()
	resp, err = handler(ctx, req) //前面的先执行； 下面的最后执行类似gin 中 next() 中间件
	fmt.Println("req:",req)
	lastTime := time.Now().Sub(now)
	fmt.Println("执行时间：",lastTime.Milliseconds())
	return
}

type TodoInfo struct{
}

func (t *TodoInfo)DoWork(ctx context.Context, td_req *pb.TodoRequest) (*pb.TodoResponse, error) {
	//接受context 参数
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		fmt.Println("metadata false")
	}
	for k,v := range md {
		fmt.Println("k:",k,"== v:",v)
	}
	fmt.Println(td_req.Todo)
	fmt.Println(td_req.Week)
	return &pb.TodoResponse{
		Done: true,
	},nil
}



func main() {
	//注册拦截器
	serviceOption := grpc.UnaryInterceptor(MyInterceptor)
	//实例化grpc
	grpcServer := grpc.NewServer(serviceOption)
	//注册服务
	pb.RegisterToDoServiceServer(grpcServer,&TodoInfo{})

	listen,err := net.Listen("tcp", "127.0.0.1:8081")
	if err != nil {
		panic(err)
	}
	err = grpcServer.Serve(listen)
	if err != nil {
		println(err)
	}


}

客户端

package main

import (
	"context"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"grpc_enum/pb"
)

//继承PerRPCCredentials  实现认证
//重写GetRequestMetadata 和 RequireTransportSecurity
type MyCredentials struct{
}

func (c *MyCredentials)GetRequestMetadata(ctx context.Context, uri ...string)(map[string]string, error){

	return map[string]string{
		"token":"123456897",
	},nil
}

func (c *MyCredentials)RequireTransportSecurity()bool{
	return false
}


func main() {
	//拦截器
	//clientInterceptor := func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error{
	//	md2 := metadata.Pairs("token","99663322poiopui")
	//	ctx = metadata.NewOutgoingContext(context.Background(), md2)
	//
	//	now := time.Now()
	//	err := invoker(ctx, method, req, reply, cc, opts ...)
	//	lastTime := time.Now().Sub(now)
	//	fmt.Println("客户端执行时间：",lastTime.Milliseconds())
	//	return err
	//}
	//opt := grpc.WithUnaryInterceptor(clientInterceptor)

	//认证
	opt := grpc.WithPerRPCCredentials(&MyCredentials{})


	conn, err := grpc.Dial("127.0.0.1:8081", grpc.WithInsecure(), opt)
	defer conn.Close()
	if err != nil {
		panic(err)
	}

	//实例化 ToDo 服务
	client := pb.NewToDoServiceClient(conn)

	md1 := metadata.New(map[string]string{
		"name" : "hello grpc",
	})

	//从过context 传递参数
	ctx := metadata.NewOutgoingContext(context.Background(), md1)

	//调用服务端DoWork方法
	rep, err := client.DoWork(ctx,&pb.TodoRequest{
		Todo:     "拦截器",
		Week:     pb.Week_Monday,
		BookMap:  map[string]string{
			"age" : "20",
		},
		//DoneTime: time.Now(),
	})

	if err != nil {
		panic(err)
	}
	fmt.Println("rep:",rep)


}