78 lines
1.7 KiB
Go
78 lines
1.7 KiB
Go
package rpc
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"crypto/x509"
|
|
"sync"
|
|
|
|
"github.com/jrick/wsrpc/v2"
|
|
)
|
|
|
|
type Client func() (*wsrpc.Client, error)
|
|
|
|
// Setup accepts RPC connection details, creates an RPC client, and returns a
|
|
// function which can be called to access the client. The returned function will
|
|
// try to handle any client disconnects by attempting to reconnect, but will
|
|
// return an error if a new connection cannot be established.
|
|
func Setup(ctx context.Context, shutdownWg *sync.WaitGroup, user, pass, addr string, cert []byte) Client {
|
|
|
|
// Create TLS options.
|
|
pool := x509.NewCertPool()
|
|
pool.AppendCertsFromPEM(cert)
|
|
tc := &tls.Config{RootCAs: pool}
|
|
tlsOpt := wsrpc.WithTLSConfig(tc)
|
|
|
|
// Create authentication options.
|
|
authOpt := wsrpc.WithBasicAuth(user, pass)
|
|
|
|
var mu sync.Mutex
|
|
var c *wsrpc.Client
|
|
|
|
// Add the graceful shutdown to the waitgroup.
|
|
shutdownWg.Add(1)
|
|
go func() {
|
|
// Wait until shutdown is signaled before shutting down.
|
|
<-ctx.Done()
|
|
|
|
if c != nil {
|
|
select {
|
|
case <-c.Done():
|
|
log.Debugf("RPC already closed (%s)", addr)
|
|
|
|
default:
|
|
log.Debugf("Closing RPC (%s)...", addr)
|
|
if err := c.Close(); err != nil {
|
|
log.Errorf("Failed to close RPC (%s): %v", addr, err)
|
|
} else {
|
|
log.Debugf("RPC closed (%s)", addr)
|
|
}
|
|
}
|
|
}
|
|
shutdownWg.Done()
|
|
}()
|
|
|
|
return func() (*wsrpc.Client, error) {
|
|
defer mu.Unlock()
|
|
mu.Lock()
|
|
|
|
if c != nil {
|
|
select {
|
|
case <-c.Done():
|
|
log.Infof("RPC client errored (%v); reconnecting...", c.Err())
|
|
c = nil
|
|
default:
|
|
return c, nil
|
|
}
|
|
}
|
|
|
|
fullAddr := "wss://" + addr + "/ws"
|
|
c, err := wsrpc.Dial(ctx, fullAddr, tlsOpt, authOpt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
log.Infof("Dialed RPC websocket %v", addr)
|
|
return c, nil
|
|
}
|
|
}
|