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

116 lines
2.4 KiB
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,
}
}