etcdb/client.go
2022-12-05 11:11:22 +01:00

123 lines
2.7 KiB
Go

package etcdb
import (
"crypto/tls"
"crypto/x509"
"flag"
"fmt"
"io/ioutil"
"log"
"strings"
etcdclient "go.etcd.io/etcd/client/v3"
"go.etcd.io/etcd/client/v3/namespace"
)
type ClientSpec struct {
URLs string
Prefix string
TLS TLSSpec
//User string
//Password string
client *etcdclient.Client
}
type TLSSpec struct {
CAPath string
KeyPath string
CertPath string
}
func (spec *ClientSpec) BindFlags(name string) {
flag.StringVar(&spec.URLs, name+"-etcd-urls", spec.URLs, "comma-separated etcd URLs ("+name+" cluster)")
flag.StringVar(&spec.Prefix, name+"-etcd-prefix", spec.Prefix, "etcd prefix ("+name+" cluster)")
// TLS flags
flag.StringVar(&spec.TLS.CAPath, name+"-ca", spec.TLS.CAPath, "etcd CA path ("+name+" cluster)")
flag.StringVar(&spec.TLS.KeyPath, name+"-key", spec.TLS.KeyPath, "etcd key path ("+name+" cluster)")
flag.StringVar(&spec.TLS.CertPath, name+"-crt", spec.TLS.CertPath, "etcd certificate path ("+name+" cluster)")
}
func (spec *ClientSpec) MustConnect() {
err := spec.Connect()
if err != nil {
log.Fatal("etcd failed to connect: ", err)
}
}
func (spec *ClientSpec) Connect() error {
cfg := etcdclient.Config{
Endpoints: strings.Split(spec.URLs, ","),
}
if caPath := spec.TLS.CAPath; caPath != "" {
if cfg.TLS == nil {
cfg.TLS = &tls.Config{}
}
caData, err := ioutil.ReadFile(caPath)
if err != nil {
return fmt.Errorf("failed to read CA file: %w", err)
}
cfg.TLS.RootCAs = x509.NewCertPool()
if !cfg.TLS.RootCAs.AppendCertsFromPEM(caData) {
return fmt.Errorf("no certificates found in %s", caPath)
}
}
if spec.TLS.KeyPath != "" || spec.TLS.CertPath != "" {
if cfg.TLS == nil {
cfg.TLS = &tls.Config{}
}
cert, err := tls.LoadX509KeyPair(spec.TLS.CertPath, spec.TLS.KeyPath)
if err != nil {
return fmt.Errorf("failed to load client key pair: %w", err)
}
cfg.TLS.Certificates = []tls.Certificate{cert}
}
client, err := etcdclient.New(cfg)
if err != nil {
return fmt.Errorf("failed to build client: %w", err)
}
prefix := spec.Prefix
if prefix != "" && !strings.HasSuffix(prefix, "/") {
prefix += "/"
}
client.KV = namespace.NewKV(client.KV, prefix)
client.Watcher = namespace.NewWatcher(client.Watcher, prefix)
client.Lease = namespace.NewLease(client.Lease, prefix)
spec.client = client
return nil
}
func (spec *ClientSpec) Client() *etcdclient.Client {
if spec.client == nil {
panic("no client, Connect() no called?")
}
client := &etcdclient.Client{}
*client = *spec.client
return client
}
func (spec *ClientSpec) Root() EtcdSpec {
return EtcdSpec{
client: spec,
}
}
func (spec *ClientSpec) Sub(prefix string) EtcdSpec {
return EtcdSpec{
client: spec,
}.Sub(prefix)
}