waf-platform/EdgeReporter/internal/rpc/rpc_client.go

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)]
}