Grpc | Token认证和自定义认证

//hello.proto
syntax = "proto3";
option go_package = "./pb;pb";
import "google/protobuf/empty.proto";

package hello;

message Login{
  string login = 1;
  string password = 2;
}

message TokenInfo {
  string token = 1;
}

message HelloWorld {
  string msg = 1;
}

service Hello {
  rpc CreateToken(Login) returns(TokenInfo) {}
  rpc SayHello(google.protobuf.Empty) returns(HelloWorld) {}
}

//jwt.go
package jwt

import (
	"context"
	"fmt"
	"pro2d/conf"
	"pro2d/utils"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"google.golang.org/grpc/metadata"
)

func CreateToken(uid string) (tokenString string) {
	token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
		"iss":      "pro2d-app-server",
		"aud":      "pro2d-app-server",
		"nbf":      time.Now().Unix(),
		"exp":      time.Now().Add(time.Hour).Unix(),
		"sub":      "pro2d",
		"uid": uid,
	})
	tokenString, err := token.SignedString([]byte(utils.Pro2DTokenSignedString))
	if err != nil {
		panic(err)
	}
	return tokenString
}

func ParseToken(tokenStr string) string {
	var clientClaims Claims
	token, err := jwt.ParseWithClaims(tokenStr, &clientClaims, func(token *jwt.Token) (interface{}, error) {
		if token.Header["alg"] != "HS256" {
			//panic("ErrInvalidAlgorithm")
			utils.Sugar.Error("ErrInvalidAlgorithm")
			return nil, nil
		}
		return []byte(utils.Pro2DTokenSignedString), nil
	})
	if err != nil {
		utils.Sugar.Error("jwt parse error")
		return ""
	}

	if !token.Valid {
		utils.Sugar.Error("ErrInvalidToken")
		return ""
	}
	return clientClaims.Uid
}


// Claims defines the struct containing the token claims.
type Claims struct {
	jwt.StandardClaims
	Uid string
}

// 从 context 的 metadata 中，取出 token
func getTokenFromContext(ctx context.Context) (string, error) {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return "", fmt.Errorf("ErrNoMetadataInContext")
	}
	// md 的类型是 type MD map[string][]string
	token, ok := md["authorization"]
	if !ok || len(token) == 0 {
		return "", fmt.Errorf("ErrNoAuthorizationInMetadata")
	}
	// 因此，token 是一个字符串数组，我们只用了 token[0]
	return token[0], nil
}

func CheckAuth(ctx context.Context) string {
	tokenStr, err := getTokenFromContext(ctx)
	if err != nil {
		utils.Sugar.Errorf("get token from context error")
		return ""
	}
	return ParseToken(tokenStr)
}

// AuthToken 自定义认证 客户端使用
type AuthToken struct {
	Token string
}

func (c AuthToken) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
	return map[string]string{
		"authorization": c.Token,
	}, nil
}

func (c AuthToken) RequireTransportSecurity() bool {
	return true
}

package main

import (
	"context"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/protobuf/types/known/emptypb"
	"log"
	"net"
	"pro2d/components/jwt"
	"pro2d/protos/pb"
)

type Server struct {
	pb.UnimplementedHelloServer
}

func (s *Server) CreateToken(ctx context.Context, in *pb.Login) (*pb.TokenInfo, error)  {
	if in.Login == "login" && in.Password == "123456" {
		return &pb.TokenInfo{Token: jwt.CreateToken(in.Login)}, nil
	}
	return nil, fmt.Errorf("login error")
}

func (s *Server) SayHello(ctx context.Context, empty *emptypb.Empty) (*pb.HelloWorld, error)  {
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		return nil, status.Errorf(codes.Unauthenticated,"ErrNoMetadataInContext")
	}
	// md 的类型是 type MD map[string][]string
	token, ok := md["authorization"]
	if !ok || len(token) == 0 {
		return nil, status.Errorf(codes.Unauthenticated,"ErrNoAuthorizationInMetadata")
	}
	login := jwt.ParseToken(token[0])
	return &pb.HelloWorld{Msg: "Hello world: " + login}, nil
}

func main() {
	// 监听本地端口
	listener, err := net.Listen("tcp", ":8948")
	if err != nil {
		log.Fatalf("net.Listen err: %v", err)
	}
	var opts []grpc.ServerOption
	// 从输入证书文件和密钥文件为服务端构造TLS凭证
	creds, err := credentials.NewServerTLSFromFile("keys/server.pem", "keys/server.key")
	if err != nil {
		log.Fatalf("Failed to generate credentials %v", err)
	}
	opts = append(opts, grpc.Creds(creds))
	// 新建gRPC服务器实例,并开启TLS认证
	grpcServer := grpc.NewServer(opts...)

	// 在gRPC服务器注册我们的服务
	pb.RegisterHelloServer(grpcServer, &Server{})
	log.Println(" net.Listing whth TLS")
	//用服务器 Serve() 方法以及我们的端口信息区实现阻塞等待，直到进程被杀死或者 Stop() 被调用
	err = grpcServer.Serve(listener)
	if err != nil {
		log.Fatalf("grpcServer.Serve err: %v", err)
	}
}

package main

import (
	"context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/protobuf/types/known/emptypb"
	"log"
	"pro2d/components/jwt"
	_ "pro2d/conf"
	"pro2d/protos/pb"
	"pro2d/utils"
)

func main() {
	var opts []grpc.DialOption
	creds, err := credentials.NewClientTLSFromFile("keys/server.pem", "pro2d")
	if err != nil {
		log.Fatal(err)
		return
	}
	opts = append(opts, grpc.WithTransportCredentials(creds))
	conn, err := grpc.Dial("localhost:8948", opts...)

	helloClient := pb.NewHelloClient(conn)
	token, err := helloClient.CreateToken(context.TODO(), &pb.Login{
		Login:    "login",
		Password: "123456",
	})
	if err != nil {
		log.Fatal(err)
		return
	}
	utils.Sugar.Debugf("token: %s", token.Token)

	opts = append(opts, grpc.WithPerRPCCredentials(&jwt.AuthToken{Token: token.Token}))
	conn2, err := grpc.Dial("localhost:8948",opts...)
	if err != nil {
		log.Fatal(err)
		return
	}

	helloClient2 := pb.NewHelloClient(conn2)
	rsp, err := helloClient2.SayHello(context.TODO(), &emptypb.Empty{})
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("sayhello rsp: %v", rsp)
}