1 files changed, 426 insertions, 0 deletions

diff --git a/secrets/vault_client.go b/secrets/vault_client.go
new file mode 100644
index 0000000..3466f48
--- /dev/null
+++ b/secrets/vault_client.go
@@ -0,0 +1,426 @@
+package secrets
+
+import (
+        "container/heap"
+        "context"
+        "encoding/json"
+        "errors"
+        "fmt"
+        "net/http"
+        "path"
+        "sync"
+        "time"
+
+        glos "code.crute.us/mcrute/golib/os"
+
+        "github.com/hashicorp/vault/api"
+        "github.com/hashicorp/vault/api/auth/approle"
+        "github.com/mitchellh/mapstructure"
+)
+
+const (
+        renewalStartPercent = 0.8
+        notificationChanLen = 100
+        defaultIncrement    = 30 * 60 // 30 minutes (as seconds)
+        defaultTimeout      = 10 * time.Second
+        renewalWindow       = 30 * time.Second
+)
+
+type vaultRenewalMinHeap []*VaultHandle
+
+func (h vaultRenewalMinHeap) Len() int           { return len(h) }
+func (h vaultRenewalMinHeap) Less(i, j int) bool { return h[i].renewAfter() < h[j].renewAfter() }
+func (h vaultRenewalMinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
+func (h *vaultRenewalMinHeap) Push(x any)        { *h = append(*h, x.(*VaultHandle)) }
+func (h vaultRenewalMinHeap) Root() *VaultHandle { return h[0] }
+
+// Convenience methods to hide the collections/heap stuff from users
+func (h *vaultRenewalMinHeap) PopHeap() *VaultHandle   { return heap.Pop(h).(*VaultHandle) }
+func (h *vaultRenewalMinHeap) PushHeap(i *VaultHandle) { heap.Push(h, i) }
+func (h *vaultRenewalMinHeap) Init()                   { heap.Init(h) }
+
+func (h *vaultRenewalMinHeap) Pop() any {
+        old := *h
+        n := len(old)
+        item := old[n-1]
+        old[n-1] = nil // avoid memory leak
+        *h = old[0 : n-1]
+        return item
+}
+
+func (h *vaultRenewalMinHeap) FindRemove(handle *VaultHandle) bool {
+        for i, hnd := range *h {
+                if hnd.name == handle.name {
+                        heap.Remove(h, i)
+                        return true
+                }
+        }
+        return false
+}
+
+type VaultHandle struct {
+        name     string
+        critical bool
+        acquired time.Time
+        secret   *api.Secret
+}
+
+var _ Handle = (*VaultHandle)(nil)
+
+func (h *VaultHandle) isAuthToken() bool {
+        return h.secret.Auth != nil
+}
+
+func (h *VaultHandle) renewAfter() time.Duration {
+        after := float64(h.leaseDuration().Nanoseconds()) * renewalStartPercent
+        afterTime := h.acquired.Add(time.Duration(after))
+        return afterTime.Sub(time.Now()).Round(time.Second)
+}
+
+func (h *VaultHandle) leaseDuration() time.Duration {
+        duration := time.Duration(h.secret.LeaseDuration) * time.Second
+        if h.isAuthToken() {
+                return time.Duration(h.secret.Auth.LeaseDuration) * time.Second
+        }
+        return duration
+}
+
+func (h *VaultHandle) renew(ctx context.Context, c VaultServiceClient, inc int) (err error) {
+        var s *api.Secret
+
+        vctx, cancel := context.WithTimeout(ctx, defaultTimeout)
+        defer cancel()
+
+        if h.isAuthToken() {
+                s, err = c.Auth().Token().RenewTokenAsSelfWithContext(vctx, h.secret.Auth.ClientToken, inc)
+                if err != nil {
+                        return err
+                }
+        } else {
+                s, err = c.Sys().RenewWithContext(vctx, h.secret.LeaseID, inc)
+                if err != nil {
+                        return err
+                }
+        }
+
+        h.secret = s
+        h.acquired = time.Now()
+
+        return nil
+}
+
+func (h *VaultHandle) Reference() string {
+        return h.name
+}
+
+type VaultServiceClient interface {
+        Auth() *api.Auth
+        Sys() *api.Sys
+        Token() string
+}
+
+type VaultClient struct {
+        sync.Mutex
+
+        client         VaultServiceClient
+        logical        *api.Logical
+        auth           *approle.AppRoleAuth
+        secrets        vaultRenewalMinHeap
+        renewIncrement int
+        notifications  chan Renewal
+}
+
+var _ Client = (*VaultClient)(nil)
+var _ ClientManager = (*VaultClient)(nil)
+
+type VaultClientConfig struct {
+        Host        string `env:"VAULT_ADDR"`
+        Token       string `env:"VAULT_TOKEN"`
+        RoleId      string `env:"VAULT_ROLE_ID"`
+        RoleSecret  string `env:"VAULT_SECRET_ID"`
+        Increment   int    `env:"VAULT_INCREMENT"`
+        AppRoleAuth *approle.AppRoleAuth
+}
+
+func (c *VaultClientConfig) Validate() error {
+        if c.Host == "" {
+                return fmt.Errorf("VaultClientConfig: Vault host is not specified")
+        }
+
+        // The presence of a token is always assumed to be valid, client errors
+        // will occur otherwise.
+        if c.Token != "" {
+                return nil
+        }
+
+        // This constructor does a bunch of validation internally so just let
+        // it do its thing and return any errors from that directly to the
+        // user.
+        ar, err := approle.NewAppRoleAuth(c.RoleId, &approle.SecretID{FromString: c.RoleSecret})
+        if err != nil {
+                return fmt.Errorf("VaultClientConfig: AppRole credentials invalid: %w", err)
+        }
+        c.AppRoleAuth = ar
+
+        return nil
+}
+
+// NewVaultClient will attempt to create a secrets.Client from the
+// passed config. Config can be nil, in which case an attempt will
+// be made to load the configuration from environment variables. See
+// VaultClientConfig for the expected names of those variables.
+func NewVaultClient(cfg *VaultClientConfig) (ClientManager, error) {
+        if cfg == nil {
+                cfg = &VaultClientConfig{}
+        }
+
+        if err := glos.UnmarshalEnvironment(cfg); err != nil {
+                return nil, err
+        }
+
+        if err := cfg.Validate(); err != nil {
+                return nil, err
+        }
+
+        vc, err := api.NewClient(api.DefaultConfig())
+        if err != nil {
+                return nil, fmt.Errorf("NewVaultClient: error building client config: %w", err)
+        }
+        vc.SetAddress(cfg.Host)
+
+        if cfg.Token != "" {
+                vc.SetToken(cfg.Token)
+        }
+
+        c := &VaultClient{
+                client:         vc,
+                logical:        vc.Logical(),
+                secrets:        vaultRenewalMinHeap{},
+                notifications:  make(chan Renewal, notificationChanLen),
+                auth:           cfg.AppRoleAuth,
+                renewIncrement: cfg.Increment,
+        }
+
+        c.secrets.Init()
+
+        if c.renewIncrement == 0 {
+                c.renewIncrement = defaultIncrement
+        }
+
+        return c, nil
+}
+
+func (c *VaultClient) Notifications() <-chan Renewal {
+        return c.notifications
+}
+
+func (c *VaultClient) Authenticate(ctx context.Context) error {
+        if c.auth == nil {
+                return c.authToken(ctx)
+        } else {
+                return c.authAppRole(ctx)
+        }
+}
+
+func (c *VaultClient) authToken(ctx context.Context) error {
+        if c.client.Token() == "" {
+                return fmt.Errorf("Authenticate: unable to authenticate, neither token nor approle provided")
+        }
+
+        vctx, cancel := context.WithTimeout(ctx, defaultTimeout)
+        defer cancel()
+
+        secret, err := c.client.Auth().Token().LookupSelfWithContext(vctx)
+        if err != nil {
+                return err
+        }
+
+        // Looking up self does not return an auth token just a map of data
+        // about the current token. Convert this into a SecretAuth so that
+        // downstream renewal code does the right thing.
+        secret.Auth = &api.SecretAuth{}
+        if err := mapstructure.Decode(secret.Data, secret.Auth); err != nil {
+                return err
+        }
+        secret.Auth.ClientToken = c.client.Token()
+
+        c.makeHandle("login", secret)
+
+        return nil
+}
+
+func (c *VaultClient) authAppRole(ctx context.Context) error {
+        vctx, cancel := context.WithTimeout(ctx, defaultTimeout)
+        defer cancel()
+
+        s, err := c.client.Auth().Login(vctx, c.auth)
+        if err != nil {
+                return fmt.Errorf("Authenticate: error logging in to vault: %w", err)
+        }
+        c.makeHandle("login", s)
+        return err
+}
+
+// makeHandle creates a secret handle and schedules it for renewal if it
+// is renewable.
+func (c *VaultClient) makeHandle(name string, s *api.Secret) Handle {
+        h := &VaultHandle{
+                name:     name,
+                critical: true, // Everything is critical unless marked otherwise
+                acquired: time.Now(),
+                secret:   s,
+        }
+
+        // If this is renewable then schedule it for renewal
+        if (s.Auth != nil && s.Auth.Renewable) || s.Renewable {
+                c.Lock()
+                c.secrets.PushHeap(h)
+                c.Unlock()
+        }
+
+        return h
+}
+
+func (c *VaultClient) read(ctx context.Context, prefix, suffix string) (Handle, error) {
+        key := path.Join(prefix, suffix)
+
+        s, err := c.logical.ReadWithContext(ctx, key)
+        if err != nil {
+                return nil, fmt.Errorf("read: error reading from Vault: %w", err)
+        }
+
+        return c.makeHandle(key, s), nil
+}
+
+func (c *VaultClient) isRecoverableDbError(err error) bool {
+        var apiErr *api.ResponseError
+        return errors.Is(api.ErrSecretNotFound, err) ||
+                (errors.As(err, &apiErr) && apiErr.StatusCode == http.StatusForbidden)
+}
+
+func (c *VaultClient) DatabaseCredential(ctx context.Context, suffix string) (*Credential, Handle, error) {
+        cred, hnd, err := c.databaseCredentialDynamic(ctx, suffix)
+        if err != nil {
+                if c.isRecoverableDbError(err) {
+                        cred, hnd, err = c.databaseCredentialStatic(ctx, suffix)
+                }
+        }
+        return cred, hnd, err
+}
+
+func (c *VaultClient) databaseCredentialDynamic(ctx context.Context, suffix string) (*Credential, Handle, error) {
+        h, err := c.read(ctx, "database/creds", suffix)
+        if err != nil {
+                return nil, nil, err
+        }
+        vh := h.(*VaultHandle)
+
+        var d Credential
+        if err = mapstructure.Decode(vh.secret.Data, &d); err != nil {
+                return nil, nil, fmt.Errorf("databaseCredentialStatic: error decoding secret: %w", err)
+        }
+
+        return &d, h, nil
+}
+
+func (c *VaultClient) databaseCredentialStatic(ctx context.Context, suffix string) (*Credential, Handle, error) {
+        h, err := c.read(ctx, "database/static-creds", suffix)
+        if err != nil {
+                return nil, nil, err
+        }
+
+        var d Credential
+        if err = mapstructure.Decode(h.(*VaultHandle).secret.Data, &d); err != nil {
+                return nil, nil, fmt.Errorf("databaseCredentialStatic: error decoding secret: %w", err)
+        }
+
+        return &d, h, nil
+}
+
+func (c *VaultClient) Secret(ctx context.Context, suffix string, out any) (Handle, error) {
+        h, err := c.read(ctx, "kv/data", suffix)
+        if err != nil {
+                return nil, err
+        }
+
+        if err = mapstructure.Decode(h.(*VaultHandle).secret.Data["data"], out); err != nil {
+                return nil, err
+        }
+
+        return h, nil
+}
+
+func (c *VaultClient) WriteSecret(ctx context.Context, suffix string, in any) error {
+        inb, err := json.Marshal(in)
+        if err != nil {
+                return fmt.Errorf("WriteSecret: error encoding json: %w", err)
+        }
+
+        if _, err = c.logical.WriteBytesWithContext(ctx, path.Join("kv/data", suffix), inb); err != nil {
+                return fmt.Errorf("WriteSecret: error writing to vault: %w", err)
+        }
+
+        return nil
+}
+
+func (c *VaultClient) Destroy(h Handle) error {
+        c.Lock()
+        defer c.Unlock()
+
+        c.secrets.FindRemove(h.(*VaultHandle))
+
+        return nil
+}
+
+func (c *VaultClient) MakeNonCritical(h Handle) error {
+        h.(*VaultHandle).critical = false
+        return nil
+}
+
+func (c *VaultClient) renewAttempt(ctx context.Context) (next time.Duration) {
+        c.Lock()
+        defer c.Unlock()
+
+        // In the absence of any other time, run once a second
+        next = 5 * time.Second
+
+        if c.secrets.Len() < 1 {
+                return
+        }
+
+        for {
+                s := c.secrets.PopHeap()
+                if s.renewAfter() < renewalWindow {
+                        // Underlying client does backoff and retry
+                        c.notifications <- Renewal{
+                                Name:     s.name,
+                                Critical: s.critical,
+                                Time:     time.Now(),
+                                Error:    s.renew(ctx, c.client, c.renewIncrement),
+                        }
+                        c.secrets.PushHeap(s)
+                } else {
+                        c.secrets.PushHeap(s)
+                        next = c.secrets.Root().renewAfter()
+                        break
+                }
+        }
+
+        return next
+}
+
+func (c *VaultClient) Run(ctx context.Context, wg *sync.WaitGroup) error {
+        wg.Add(1)
+        defer wg.Done()
+
+        for {
+                sleepTime := c.renewAttempt(ctx)
+
+                select {
+                case <-time.After(sleepTime):
+                        continue
+                case <-ctx.Done():
+                        return nil
+                }
+        }
+}