Harden Vault DB lease handling (#65)

* Renew and track Vault DB leases

* Revoke Vault DB leases on renewal failure

Author: haasonsaas

internal/connectors/vaultdb/client.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net/http"
 	"strings"
 	"time"
@@ -19,6 +20,7 @@ type HTTPClientConfig struct {
 	Token       string
 	Namespace   string
 	Client      *http.Client
+	RenewRetry  resilience.RetryConfig
 	RevokeRetry resilience.RetryConfig
 }
 
@@ -27,6 +29,7 @@ type HTTPClient struct {
 	token       string
 	namespace   string
 	client      *http.Client
+	renewRetry  resilience.RetryConfig
 	revokeRetry resilience.RetryConfig
 }
 
@@ -45,11 +48,22 @@ func NewHTTPClient(cfg HTTPClientConfig) *HTTPClient {
 	if revokeRetry.MaxDelay == 0 {
 		revokeRetry.MaxDelay = time.Second
 	}
+	renewRetry := cfg.RenewRetry
+	if renewRetry.MaxAttempts == 0 {
+		renewRetry.MaxAttempts = 3
+	}
+	if renewRetry.InitialDelay == 0 {
+		renewRetry.InitialDelay = 100 * time.Millisecond
+	}
+	if renewRetry.MaxDelay == 0 {
+		renewRetry.MaxDelay = time.Second
+	}
 	return &HTTPClient{
 		baseURL:     strings.TrimRight(cfg.BaseURL, "/"),
 		token:       cfg.Token,
 		namespace:   cfg.Namespace,
 		client:      client,
+		renewRetry:  renewRetry,
 		revokeRetry: revokeRetry,
 	}
 }
@@ -79,6 +93,7 @@ func (c *HTTPClient) GenerateCredentials(ctx context.Context, role string) (*Lea
 	var payload struct {
 		LeaseID       string `json:"lease_id"`
 		LeaseDuration int64  `json:"lease_duration"`
+		Renewable     bool   `json:"renewable"`
 		Data          struct {
 			Username string `json:"username"`
 			Password string `json:"password"`
@@ -92,15 +107,71 @@ func (c *HTTPClient) GenerateCredentials(ctx context.Context, role string) (*Lea
 		Password:      payload.Data.Password,
 		LeaseID:       payload.LeaseID,
 		LeaseDuration: time.Duration(payload.LeaseDuration) * time.Second,
+		Renewable:     payload.Renewable,
 	}, nil
 }
 
+func (c *HTTPClient) RenewLease(ctx context.Context, leaseID string, increment time.Duration) (*LeaseCredentials, error) {
+	return resilience.RetryValue(ctx, c.renewRetry, func(ctx context.Context) (*LeaseCredentials, error) {
+		return c.renewLeaseOnce(ctx, leaseID, increment)
+	})
+}
+
 func (c *HTTPClient) RevokeLease(ctx context.Context, leaseID string) error {
 	return resilience.Retry(ctx, c.revokeRetry, func(ctx context.Context) error {
 		return c.revokeLeaseOnce(ctx, leaseID)
 	})
 }
 
+func (c *HTTPClient) renewLeaseOnce(ctx context.Context, leaseID string, increment time.Duration) (*LeaseCredentials, error) {
+	payload, err := json.Marshal(map[string]any{
+		"lease_id":  leaseID,
+		"increment": int(math.Ceil(increment.Seconds())),
+	})
+	if err != nil {
+		return nil, resilience.Permanent(err)
+	}
+	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/sys/leases/renew", bytes.NewReader(payload))
+	if err != nil {
+		return nil, resilience.Permanent(err)
+	}
+	c.applyHeaders(request)
+	request.Header.Set("Content-Type", "application/json")
+
+	response, err := c.client.Do(request)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		_ = response.Body.Close()
+	}()
+	body, err := io.ReadAll(response.Body)
+	if err != nil {
+		return nil, err
+	}
+	if response.StatusCode >= 400 {
+		wrapped := fmt.Errorf("%w: vault renew lease returned %d: %s", core.ErrForbidden, response.StatusCode, string(body))
+		if response.StatusCode == http.StatusTooManyRequests || response.StatusCode >= http.StatusInternalServerError {
+			return nil, wrapped
+		}
+		return nil, resilience.Permanent(wrapped)
+	}
+
+	var renewed struct {
+		LeaseID       string `json:"lease_id"`
+		LeaseDuration int64  `json:"lease_duration"`
+		Renewable     bool   `json:"renewable"`
+	}
+	if err := json.Unmarshal(body, &renewed); err != nil {
+		return nil, resilience.Permanent(err)
+	}
+	return &LeaseCredentials{
+		LeaseID:       renewed.LeaseID,
+		LeaseDuration: time.Duration(renewed.LeaseDuration) * time.Second,
+		Renewable:     renewed.Renewable,
+	}, nil
+}
+
 func (c *HTTPClient) revokeLeaseOnce(ctx context.Context, leaseID string) error {
 	body, err := json.Marshal(map[string]string{"lease_id": leaseID})
 	if err != nil {

internal/connectors/vaultdb/client_test.go

@@ -25,8 +25,19 @@ func TestHTTPClient_GenerateCredentialsAndRevokeLease(t *testing.T) {
 			_, _ = w.Write([]byte(`{
 				"lease_id":"database/creds/analytics_ro/123",
 				"lease_duration":600,
+				"renewable":true,
 				"data":{"username":"dyn-user","password":"dyn-pass"}
 			}`))
+		case "/v1/sys/leases/renew":
+			body, _ := io.ReadAll(r.Body)
+			if string(body) != `{"increment":1800,"lease_id":"database/creds/analytics_ro/123"}` {
+				t.Fatalf("renew body = %s, want lease renew payload", string(body))
+			}
+			_, _ = w.Write([]byte(`{
+				"lease_id":"database/creds/analytics_ro/123",
+				"lease_duration":1800,
+				"renewable":true
+			}`))
 		case "/v1/sys/leases/revoke":
 			body, _ := io.ReadAll(r.Body)
 			if string(body) != `{"lease_id":"database/creds/analytics_ro/123"}` {
@@ -51,6 +62,13 @@ func TestHTTPClient_GenerateCredentialsAndRevokeLease(t *testing.T) {
 	if lease.Username != "dyn-user" || lease.LeaseID != "database/creds/analytics_ro/123" {
 		t.Fatalf("lease = %#v, want parsed vault lease", lease)
 	}
+	renewed, err := client.RenewLease(context.Background(), lease.LeaseID, 30*time.Minute)
+	if err != nil {
+		t.Fatalf("RenewLease() error = %v", err)
+	}
+	if renewed.LeaseDuration != 30*time.Minute {
+		t.Fatalf("renewed lease = %#v, want renewed duration", renewed)
+	}
 	if err := client.RevokeLease(context.Background(), lease.LeaseID); err != nil {
 		t.Fatalf("RevokeLease() error = %v", err)
 	}

internal/connectors/vaultdb/connector.go

@@ -16,10 +16,12 @@ type LeaseCredentials struct {
 	Password      string
 	LeaseID       string
 	LeaseDuration time.Duration
+	Renewable     bool
 }
 
 type Client interface {
 	GenerateCredentials(ctx context.Context, role string) (*LeaseCredentials, error)
+	RenewLease(ctx context.Context, leaseID string, increment time.Duration) (*LeaseCredentials, error)
 	RevokeLease(ctx context.Context, leaseID string) error
 }
 
@@ -117,6 +119,14 @@ func (c *Connector) Issue(ctx context.Context, req core.IssueRequest) (*core.Iss
 	if err != nil {
 		return nil, err
 	}
+	extendedLease, leaseExpiresAt, err := c.extendLeaseForGrant(ctx, lease, req.Grant.ExpiresAt)
+	if err != nil {
+		if lease != nil && lease.LeaseID != "" {
+			_ = c.client.RevokeLease(ctx, lease.LeaseID)
+		}
+		return nil, err
+	}
+	lease = extendedLease
 	dsn, err := renderDSN(c.roleDSNs[req.Resource.Name], lease)
 	if err != nil {
 		return nil, err
@@ -125,16 +135,17 @@ func (c *Connector) Issue(ctx context.Context, req core.IssueRequest) (*core.Iss
 	return &core.IssuedArtifact{
 		Kind: core.ArtifactKindWrappedSecret,
 		Metadata: map[string]string{
-			"artifact_id": "art_" + req.Grant.ID,
-			"lease_id":    lease.LeaseID,
-			"db_role":     req.Resource.Name,
+			"artifact_id":      "art_" + req.Grant.ID,
+			"lease_id":         lease.LeaseID,
+			"lease_expires_at": leaseExpiresAt.UTC().Format(time.RFC3339),
+			"db_role":          req.Resource.Name,
 		},
 		SecretData: map[string]string{
 			"username": lease.Username,
 			"password": lease.Password,
 			"dsn":      dsn,
 		},
-		ExpiresAt: minTime(req.Grant.ExpiresAt, time.Now().Add(lease.LeaseDuration)),
+		ExpiresAt: minTime(req.Grant.ExpiresAt, leaseExpiresAt),
 	}, nil
 }
 
@@ -203,6 +214,50 @@ func renderDSN(renderPattern string, lease *LeaseCredentials) (string, error) {
 	return builder.String(), nil
 }
 
+func (c *Connector) extendLeaseForGrant(ctx context.Context, lease *LeaseCredentials, grantExpiresAt time.Time) (*LeaseCredentials, time.Time, error) {
+	now := time.Now()
+	leaseExpiresAt := now.Add(lease.LeaseDuration)
+	if grantExpiresAt.IsZero() || !grantExpiresAt.After(leaseExpiresAt) {
+		return lease, leaseExpiresAt, nil
+	}
+	if lease.LeaseID == "" || !lease.Renewable {
+		return nil, time.Time{}, fmt.Errorf("%w: vault lease expires before the requested grant ttl and is not renewable", core.ErrForbidden)
+	}
+
+	remaining := time.Until(grantExpiresAt)
+	renewed, err := c.client.RenewLease(ctx, lease.LeaseID, remaining)
+	if err != nil {
+		return nil, time.Time{}, err
+	}
+	merged := mergeLeaseCredentials(lease, renewed)
+	leaseExpiresAt = time.Now().Add(merged.LeaseDuration)
+	if grantExpiresAt.After(leaseExpiresAt) {
+		return nil, time.Time{}, fmt.Errorf("%w: vault lease renewal for %q was capped before the requested grant ttl", core.ErrForbidden, lease.LeaseID)
+	}
+	return merged, leaseExpiresAt, nil
+}
+
+func mergeLeaseCredentials(current *LeaseCredentials, updated *LeaseCredentials) *LeaseCredentials {
+	if updated == nil {
+		return current
+	}
+	merged := *current
+	if updated.Username != "" {
+		merged.Username = updated.Username
+	}
+	if updated.Password != "" {
+		merged.Password = updated.Password
+	}
+	if updated.LeaseID != "" {
+		merged.LeaseID = updated.LeaseID
+	}
+	if updated.LeaseDuration > 0 {
+		merged.LeaseDuration = updated.LeaseDuration
+	}
+	merged.Renewable = updated.Renewable
+	return &merged
+}
+
 func minTime(a time.Time, b time.Time) time.Time {
 	if a.IsZero() {
 		return b