1.4.5.2

EdgeReporter/internal/rpc/rpc_client.go

@@ -0,0 +1,191 @@
+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)]
+}

EdgeReporter/internal/rpc/rpc_client_test.go

@@ -0,0 +1,5 @@
+package rpc
+
+import (
+	_ "github.com/iwind/TeaGo/bootstrap"
+)

EdgeReporter/internal/rpc/rpc_utils.go

@@ -0,0 +1,30 @@
+package rpc
+
+import (
+	"github.com/TeaOSLab/EdgeReporter/internal/configs"
+	"sync"
+)
+
+var sharedRPC *RPCClient = nil
+var locker = &sync.Mutex{}
+
+func SharedRPC() (*RPCClient, error) {
+	locker.Lock()
+	defer locker.Unlock()
+
+	if sharedRPC != nil {
+		return sharedRPC, nil
+	}
+
+	config, err := configs.LoadAPIConfig()
+	if err != nil {
+		return nil, err
+	}
+	client, err := NewRPCClient(config)
+	if err != nil {
+		return nil, err
+	}
+
+	sharedRPC = client
+	return sharedRPC, nil
+}