JCS-pub/common/pkgs/rpc/coordinator/pool.go

package corrpc

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"os"

	"gitlink.org.cn/cloudream/common/consts/errorcode"
	"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type PoolConfig struct {
	Address string
	Conn    rpc.PoolConfig
}

type PoolConfigJSON struct {
	Address    string `json:"address"`
	RootCA     string `json:"rootCA"`
	ClientCert string `json:"clientCert"`
	ClientKey  string `json:"clientKey"`
}

func (c *PoolConfigJSON) Build(tokenProv rpc.AccessTokenProvider) (*PoolConfig, error) {
	pc := &PoolConfig{
		Address: c.Address,
	}
	pc.Conn.AccessTokenProvider = tokenProv

	rootCA, err := os.ReadFile(c.RootCA)
	if err != nil {
		return nil, fmt.Errorf("load root ca: %v", err)
	}
	pc.Conn.RootCA = x509.NewCertPool()
	if !pc.Conn.RootCA.AppendCertsFromPEM(rootCA) {
		return nil, fmt.Errorf("failed to parse root ca")
	}

	if c.ClientCert != "" && c.ClientKey != "" {
		cert, err := tls.LoadX509KeyPair(c.ClientCert, c.ClientKey)
		if err != nil {
			return nil, fmt.Errorf("load client cert: %v", err)
		}
		pc.Conn.ClientCert = &cert
	} else if tokenProv == nil {
		return nil, fmt.Errorf("must provide client cert or access token provider")
	}

	return pc, nil
}

func (c *PoolConfigJSON) BuildTempClient() (*TempClient, error) {
	rootCA, err := os.ReadFile(c.RootCA)
	if err != nil {
		return nil, fmt.Errorf("load root ca: %v", err)
	}
	rootCAs := x509.NewCertPool()
	if !rootCAs.AppendCertsFromPEM(rootCA) {
		return nil, fmt.Errorf("failed to parse root ca")
	}

	gcon, err := grpc.NewClient(c.Address,
		grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{
			RootCAs:    rootCAs,
			ServerName: rpc.ClientAPISNIV1,
			NextProtos: []string{"h2"},
		})),
	)
	if err != nil {
		return nil, err
	}

	return &TempClient{
		Client: Client{
			con:      gcon,
			cli:      NewCoordinatorClient(gcon),
			pool:     nil,
			fusedErr: nil,
		},
	}, nil
}

type Pool struct {
	cfg      PoolConfig
	connPool *rpc.ConnPool
}

func NewPool(cfg PoolConfig) *Pool {
	return &Pool{
		cfg:      cfg,
		connPool: rpc.NewConnPool(cfg.Conn),
	}
}

func (p *Pool) Get() *Client {
	con, err := p.connPool.GetConnection(p.cfg.Address)
	if err != nil {
		return &Client{
			con:      nil,
			cli:      nil,
			pool:     p,
			fusedErr: rpc.Failed(errorcode.OperationFailed, err.Error()),
		}
	}

	return &Client{
		con:      con,
		cli:      NewCoordinatorClient(con),
		pool:     p,
		fusedErr: nil,
	}
}