package main

import (
    "flag"
    "fmt"

    "grpc-common/ucenter/types/login"
    "grpc-common/ucenter/types/register"
    "ucenter/internal/config"
    "ucenter/internal/middleware"
    "ucenter/internal/server"
    "ucenter/internal/svc"

    "github.com/zeromicro/go-zero/core/conf"
    "github.com/zeromicro/go-zero/core/service"
    "github.com/zeromicro/go-zero/zrpc"
    "google.golang.org/grpc"
    "google.golang.org/grpc/reflection"
)

var configFile = flag.String("f", "etc/config.yaml", "the config file")

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) {
        register.RegisterRegisterServer(grpcServer, server.NewRegisterServer(ctx))
        login.RegisterLoginServer(grpcServer, server.NewLoginServer(ctx))
        if c.Mode == service.DevMode || c.Mode == service.TestMode {
            reflection.Register(grpcServer)
        }
    })
    defer s.Stop()

    s.AddUnaryInterceptors(middleware.TestRpcMiddleware)

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

package middleware

import (
    "context"
    "fmt"

    "google.golang.org/grpc"
)

// rpc-server服务端拦截器
func TestRpcMiddleware(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
    fmt.Println("rpc-server调用")
    fmt.Printf("req =========> %+v \n", req)
    fmt.Printf("info =========> %+v \n", info)
    fmt.Println("rpc-server调用结束")
    return handler(ctx, req)
}

package svc

import (
    "grpc-common/ucenter/uc_client"
    "ucenter-api/internal/config"
    "ucenter-api/internal/middleware"

    "github.com/zeromicro/go-zero/zrpc"
)

type ServiceContext struct {
    Config        config.Config
    UCRegisterRpc uc_client.Register
    UCLoginRpc    uc_client.Login
}

func NewServiceContext(c config.Config) *ServiceContext {
    return &ServiceContext{
        Config:        c,
        UCRegisterRpc: uc_client.NewRegister(zrpc.MustNewClient(c.MyGrpc)),
        UCLoginRpc:    uc_client.NewLogin(zrpc.MustNewClient(c.MyGrpc, zrpc.WithUnaryClientInterceptor(middleware.TestClientRpcMiddleware))),
    }
}

package middleware

import (
    "context"
    "fmt"

    "google.golang.org/grpc"
)

// rpc-client客户端拦截器
func TestClientRpcMiddleware(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    fmt.Println("rpc-client调用前")
    err := invoker(ctx, method, req, reply, cc, opts...)
    if err != nil {
        return err
    }
    fmt.Println("rpc-client调用后")
    return nil
}

package middleware

import (
    "context"
    "fmt"

    "google.golang.org/grpc"
    "google.golang.org/grpc/metadata"
)

// rpc-client客户端拦截器
func TestClientRpcMiddleware(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
    fmt.Println("rpc-client调用前")

    // 通过上下文传值
    md := metadata.New(map[string]string{
        "user": "chf",
    })
    ctx = metadata.NewOutgoingContext(ctx, md)

    err := invoker(ctx, method, req, reply, cc, opts...)
    if err != nil {
        return err
    }
    fmt.Println("rpc-client调用后")
    return nil
}

// 定义一个 SayHello 一元 rpc 方法，请求体和响应体必填。
func (l *LoginLogic) Login(in *login.LoginReq) (*login.LoginRes, error) {
    if md, ok := metadata.FromIncomingContext(l.ctx); ok {
        tmp := md.Get("user") // 返回切片
        if len(tmp) > 0 {
            fmt.Printf("metadata-tmp: %v+ \n", tmp[0])
        }
    }

    // 校验手机号
    member, err := l.MemberDomain.FindByPhone(context.Background(), in.Phone)
    if err != nil {
        return nil, errors.New("登录失败")
    }
    if member == nil {
        return nil, errors.New("此用户未注册")
    }
    token, err := tools.GenToken(member.Id, member.Name)
    if err != nil {
        return nil, err
    }
    return &login.LoginRes{
        Token: token,
        Name:  member.Name,
        Phone: member.Phone,
        Id:    member.Id,
    }, nil
}