package rpc import ( "context" "crypto/tls" "encoding/base64" "errors" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/TeaOSLab/EdgeReporter/internal/configs" teaconst "github.com/TeaOSLab/EdgeReporter/internal/const" "github.com/TeaOSLab/EdgeReporter/internal/encrypt" "github.com/TeaOSLab/EdgeReporter/internal/utils" "github.com/iwind/TeaGo/maps" "github.com/iwind/TeaGo/rands" "google.golang.org/grpc" "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/encoding/gzip" "google.golang.org/grpc/keepalive" "google.golang.org/grpc/metadata" "net/url" "sync" "time" ) // RPCClient RPC客户端 type RPCClient struct { apiConfig *configs.APIConfig conns []*grpc.ClientConn locker sync.Mutex } // NewRPCClient 构造新的RPC客户端 func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { if apiConfig == nil { return nil, errors.New("api config should not be nil") } client := &RPCClient{ apiConfig: apiConfig, } err := client.init() if err != nil { return nil, err } return client, nil } func (this *RPCClient) SysSettingRPC() pb.SysSettingServiceClient { return pb.NewSysSettingServiceClient(this.pickConn()) } func (this *RPCClient) SysLockerRPC() pb.SysLockerServiceClient { return pb.NewSysLockerServiceClient(this.pickConn()) } func (this *RPCClient) NodeLogRPC() pb.NodeLogServiceClient { return pb.NewNodeLogServiceClient(this.pickConn()) } func (this *RPCClient) APINodeRPC() pb.APINodeServiceClient { return pb.NewAPINodeServiceClient(this.pickConn()) } func (this *RPCClient) ReportNodeRPC() pb.ReportNodeServiceClient { return pb.NewReportNodeServiceClient(this.pickConn()) } func (this *RPCClient) ReportResultRPC() pb.ReportResultServiceClient { return pb.NewReportResultServiceClient(this.pickConn()) } // Context 构造用户上下文 func (this *RPCClient) Context() context.Context { ctx := context.Background() m := maps.Map{ "timestamp": time.Now().Unix(), "type": "report", "userId": 0, } method, err := encrypt.NewMethodInstance(teaconst.EncryptMethod, this.apiConfig.Secret, this.apiConfig.NodeId) if err != nil { utils.PrintError(err) return context.Background() } data, err := method.Encrypt(m.AsJSON()) if err != nil { utils.PrintError(err) return context.Background() } token := base64.StdEncoding.EncodeToString(data) ctx = metadata.AppendToOutgoingContext(ctx, "nodeId", this.apiConfig.NodeId, "token", token) return ctx } // UpdateConfig 修改配置 func (this *RPCClient) UpdateConfig(config *configs.APIConfig) error { this.apiConfig = config this.locker.Lock() err := this.init() this.locker.Unlock() return err } // 初始化 func (this *RPCClient) init() error { // 重新连接 var conns = []*grpc.ClientConn{} for _, endpoint := range this.apiConfig.RPCEndpoints { u, err := url.Parse(endpoint) if err != nil { return errors.New("parse endpoint failed: " + err.Error()) } var conn *grpc.ClientConn var callOptions = grpc.WithDefaultCallOptions( grpc.MaxCallRecvMsgSize(128*1024*1024), grpc.UseCompressor(gzip.Name), ) var keepaliveParams = grpc.WithKeepaliveParams(keepalive.ClientParameters{ Time: 30 * time.Second, }) if u.Scheme == "http" { conn, err = grpc.Dial(u.Host, grpc.WithTransportCredentials(insecure.NewCredentials()), callOptions, keepaliveParams) } else if u.Scheme == "https" { conn, err = grpc.Dial(u.Host, grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{ InsecureSkipVerify: true, })), callOptions, keepaliveParams) } else { return errors.New("parse endpoint failed: invalid scheme '" + u.Scheme + "'") } if err != nil { return err } conns = append(conns, conn) } if len(conns) == 0 { return errors.New("[RPC]no available endpoints") } // 这里不需要加锁,因为会和pickConn冲突 this.conns = conns return nil } // 随机选择一个连接 func (this *RPCClient) pickConn() *grpc.ClientConn { this.locker.Lock() defer this.locker.Unlock() // 检查连接状态 var countConns = len(this.conns) if countConns > 0 { if countConns == 1 { return this.conns[0] } for _, state := range []connectivity.State{ connectivity.Ready, connectivity.Idle, connectivity.Connecting, connectivity.TransientFailure, } { var availableConns = []*grpc.ClientConn{} for _, conn := range this.conns { if conn.GetState() == state { availableConns = append(availableConns, conn) } } if len(availableConns) > 0 { return this.randConn(availableConns) } } } return this.randConn(this.conns) } func (this *RPCClient) randConn(conns []*grpc.ClientConn) *grpc.ClientConn { var l = len(conns) if l == 0 { return nil } if l == 1 { return conns[0] } return conns[rands.Int(0, l-1)] }