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