Merge pull request #12 from kiwicom/vault-token-cache

feat(cache): SA service accoutns + vault tokens

controllers/vaultsecret_controller.go

@@ -21,6 +21,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"sync"
 	"time"
 
 	vaultAPI "github.com/hashicorp/vault/api"
@@ -51,6 +52,8 @@ type VaultSecretReconciler struct {
 	VaultConfig   vault.AppConfig
 	VaultClient   *vaultAPI.Client
 	K8ClientSet   *kubernetes.Clientset
+	authSACache   map[string]*vault.AuthServiceAccount
+	saCacheMx     sync.RWMutex
 }
 
 func NewVaultReconciler(mgr manager.Manager, cfg vault.AppConfig, vaultClient *vaultAPI.Client,
@@ -62,6 +65,7 @@ func NewVaultReconciler(mgr manager.Manager, cfg vault.AppConfig, vaultClient *v
 		VaultConfig:   cfg,
 		VaultClient:   vaultClient,
 		K8ClientSet:   k8ClientSet,
+		authSACache:   make(map[string]*vault.AuthServiceAccount),
 	}
 	err := reconciler.setupWithManager(mgr)
 	if err != nil {
@@ -132,17 +136,12 @@ func (r *VaultSecretReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	if vaultSecret.Spec.Auth.Token != "" {
 		tokener = vault.NewAuthToken(vaultSecret.Spec.Auth.Token)
 	} else {
-		saRef := vaultSecret.Spec.Auth.ServiceAccountRef
-		vaultClient, err := r.VaultClient.Clone()
+		saAccount, err := r.getAuthServiceAccount(vaultSecret)
 		if err != nil {
-			return ctrl.Result{}, fmt.Errorf("vault client clone: %w", err)
+			return ctrl.Result{}, fmt.Errorf("get auth service account: %w", err)
 		}
-		if err := vaultClient.SetAddress(vaultSecret.Spec.Addr); err != nil {
-			return ctrl.Result{}, fmt.Errorf("vault set address: %w", err)
-		}
-		tokener = vault.NewAuthServiceAccount(vaultClient, r.K8ClientSet, saRef.Name, vaultSecret.Namespace, saRef.Role, saRef.AuthPath, false)
+		tokener = saAccount
 	}
-
 	reader, err := vault.NewReader(tokener, &vaultSecret, logger, &r.VaultConfig)
 	if err != nil {
 		r.EventRecorder.Warning(&vaultSecret, "vault failed", err)
@@ -205,6 +204,34 @@ func (r *VaultSecretReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 	return ctrl.Result{RequeueAfter: reconcileAfter}, nil
 }
 
+func (r *VaultSecretReconciler) cachedSA(id string) *vault.AuthServiceAccount {
+	r.saCacheMx.Lock()
+	defer r.saCacheMx.Unlock()
+	return r.authSACache[id]
+}
+
+func (r *VaultSecretReconciler) getAuthServiceAccount(vaultSecret k8skiwicomv1.VaultSecret) (*vault.AuthServiceAccount, error) {
+	saRef := vaultSecret.Spec.Auth.ServiceAccountRef
+	id := fmt.Sprintf("%s-%s-%s-%s-%s", vaultSecret.Spec.Addr, vaultSecret.Namespace, saRef.Name, saRef.Role, saRef.AuthPath)
+	saAccount := r.cachedSA(id)
+	if saAccount != nil {
+		return saAccount, nil
+	}
+
+	vaultClient, err := r.VaultClient.Clone()
+	if err != nil {
+		return nil, fmt.Errorf("vault client clone: %w", err)
+	}
+	if err := vaultClient.SetAddress(vaultSecret.Spec.Addr); err != nil {
+		return nil, fmt.Errorf("vault set address: %w", err)
+	}
+	saAccount = vault.NewAuthServiceAccount(vaultClient, r.K8ClientSet, saRef.Name, vaultSecret.Namespace, saRef.Role, saRef.AuthPath, false)
+	r.saCacheMx.Lock()
+	defer r.saCacheMx.Unlock()
+	r.authSACache[id] = saAccount
+	return saAccount, nil
+}
+
 // setupWithManager sets up the controller with the Manager.
 func (r *VaultSecretReconciler) setupWithManager(mgr ctrl.Manager) error {
 	return ctrl.NewControllerManagedBy(mgr).

pkg/vault/auth.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"sync"
+	"time"
 
 	vaultApi "github.com/hashicorp/vault/api"
 	authenticationv1 "k8s.io/api/authentication/v1"
@@ -28,18 +30,21 @@ func (a AuthToken) Token() (string, error) {
 }
 
 type AuthServiceAccount struct {
-	name        string
-	namespace   string
-	role        string
-	path        string
-	vaultClient *vaultApi.Client
-	autoMount   bool
-	k8ClientSet *kubernetes.Clientset
+	name             string
+	namespace        string
+	role             string
+	path             string
+	vaultClient      *vaultApi.Client
+	autoMount        bool
+	k8ClientSet      *kubernetes.Clientset
+	cacheMx          sync.RWMutex
+	cachedVaultToken string
+	vaultTokenExpire time.Time
 }
 
 func NewAuthServiceAccount(vaultClient *vaultApi.Client, k8ClientSet *kubernetes.Clientset,
-	name, namespace, role, path string, automount bool) AuthServiceAccount {
-	return AuthServiceAccount{
+	name, namespace, role, path string, automount bool) *AuthServiceAccount {
+	return &AuthServiceAccount{
 		name:        name,
 		namespace:   namespace,
 		role:        role,
@@ -49,8 +54,18 @@ func NewAuthServiceAccount(vaultClient *vaultApi.Client, k8ClientSet *kubernetes
 		k8ClientSet: k8ClientSet,
 	}
 }
+func (a *AuthServiceAccount) cachedToken() string {
+	a.cacheMx.Lock()
+	defer a.cacheMx.Unlock()
+	return a.cachedVaultToken
+}
+
+func (a *AuthServiceAccount) Token() (string, error) {
+	vaultToken := a.cachedToken()
+	if vaultToken != "" && time.Now().Add(30*time.Second).Before(a.vaultTokenExpire) {
+		return vaultToken, nil
+	}
 
-func (a AuthServiceAccount) Token() (string, error) {
 	jwtToken, err := a.fetchJWT()
 	if err != nil {
 		return "", fmt.Errorf("could not fetch JWT token: %w", err)
@@ -61,6 +76,7 @@ func (a AuthServiceAccount) Token() (string, error) {
 		"jwt":  jwtToken,
 	}
 
+	t := time.Now()
 	resp, err := a.vaultClient.Logical().Write(a.path, data)
 	if err != nil {
 		return "", fmt.Errorf("failed to login to Vault with JWT: %w", err)
@@ -71,14 +87,21 @@ func (a AuthServiceAccount) Token() (string, error) {
 		return "", fmt.Errorf("could not read auth token from response: %w", err)
 	}
 
-	// resp has resp.LeaseDuration and resp.Tokener.LeaseDuration which can cause confusion
-	// with autocomplete, use TokenTTL method instead
-	// ttl, err := resp.TokenTTL()
+	duration, err := resp.TokenTTL()
+	if err != nil {
+		return "", fmt.Errorf("could not read auth token TTL from response: %w", err)
+	}
+	expireIn := t.Add(duration)
+
+	a.cacheMx.Lock()
+	defer a.cacheMx.Unlock()
+	a.cachedVaultToken = token
+	a.vaultTokenExpire = expireIn
 
 	return token, nil
 }
 
-func (a AuthServiceAccount) fetchJWT() (string, error) {
+func (a *AuthServiceAccount) fetchJWT() (string, error) {
 	if a.autoMount {
 		return fetchJWTFromAutoMountedSecret()
 	}