120 lines
2.5 KiB
Go
120 lines
2.5 KiB
Go
package hubrpc
|
|
|
|
import (
|
|
"crypto/tls"
|
|
"crypto/x509"
|
|
"fmt"
|
|
"os"
|
|
|
|
"gitlink.org.cn/cloudream/common/consts/errorcode"
|
|
"gitlink.org.cn/cloudream/jcs-pub/common/pkgs/rpc"
|
|
jcstypes "gitlink.org.cn/cloudream/jcs-pub/common/types"
|
|
)
|
|
|
|
type PoolConfig struct {
|
|
Conn rpc.PoolConfig
|
|
AddrProvider func(id jcstypes.HubID) string
|
|
}
|
|
|
|
type PoolConfigJSON struct {
|
|
RootCA string `json:"rootCA"`
|
|
ClientCert string `json:"clientCert"`
|
|
ClientKey string `json:"clientKey"`
|
|
}
|
|
|
|
func (c *PoolConfigJSON) Build(tokenProv rpc.AccessTokenProvider) (*PoolConfig, error) {
|
|
pc := &PoolConfig{}
|
|
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
|
|
}
|
|
|
|
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(ip string, port int) *Client {
|
|
addr := fmt.Sprintf("%s:%d", ip, port)
|
|
con, err := p.connPool.GetConnection(addr)
|
|
if err != nil {
|
|
return &Client{
|
|
addr: addr,
|
|
con: nil,
|
|
pool: p,
|
|
fusedErr: rpc.Failed(errorcode.OperationFailed, err.Error()),
|
|
}
|
|
}
|
|
|
|
return &Client{
|
|
addr: addr,
|
|
con: con,
|
|
cli: NewHubClient(con),
|
|
pool: p,
|
|
}
|
|
}
|
|
|
|
func (p *Pool) GetByID(hubID jcstypes.HubID) *Client {
|
|
if p.cfg.AddrProvider == nil {
|
|
return &Client{
|
|
addr: "",
|
|
con: nil,
|
|
pool: p,
|
|
fusedErr: rpc.Failed(errorcode.OperationFailed, "no address provider"),
|
|
}
|
|
}
|
|
|
|
addr := p.cfg.AddrProvider(hubID)
|
|
if addr == "" {
|
|
return &Client{
|
|
addr: "",
|
|
con: nil,
|
|
pool: p,
|
|
fusedErr: rpc.Failed(errorcode.OperationFailed, "no address for hub %v", hubID),
|
|
}
|
|
}
|
|
|
|
con, err := p.connPool.GetConnection(addr)
|
|
if err != nil {
|
|
return &Client{
|
|
addr: addr,
|
|
con: nil,
|
|
pool: p,
|
|
fusedErr: rpc.Failed(errorcode.OperationFailed, err.Error()),
|
|
}
|
|
}
|
|
|
|
return &Client{
|
|
addr: addr,
|
|
con: con,
|
|
cli: NewHubClient(con),
|
|
pool: p,
|
|
}
|
|
}
|