Enable persistent agent caching (hashicorp#229)

Controlled by the "vault.hashicorp.com/agent-cache-enable" annotation, with a few caveats: If "vault.hashicorp.com/agent-cache-enable": "true", ...and if init and sidecar enabled, persistent caching enabled ...and if sidecar only, just memory caching without persistence enabled ...and if init only, no caching at all When persistent caching is enabled, a `vault-agent-cache` memory volume is mounted at `/vault/agent-cache` on the init and sidecar containers.
RemcoBuddelmeijer · Mar 16, 2021 · 7c993e9 · 7c993e9
1 parent 3db9da0
commit 7c993e9
Show file tree

Hide file tree

Showing 8 changed files with 339 additions and 12 deletions.
diff --git a/agent-inject/agent/agent.go b/agent-inject/agent/agent.go
@@ -27,6 +27,7 @@ const (
 	DefaultAgentCacheEnable              = "false"
 	DefaultAgentCacheUseAutoAuthToken    = "true"
 	DefaultAgentCacheListenerPort        = "8200"
+	DefaultAgentCacheExitOnErr           = false
 	DefaultAgentUseLeaderElector         = false
 )
 
@@ -230,6 +231,13 @@ type VaultAgentCache struct {
 
 	// UseAutoAuthToken configures whether the auto auth token is used in cache requests
 	UseAutoAuthToken string
+
+	// Persist marks whether persistent caching is enabled or not
+	Persist bool
+
+	// ExitOnErr configures whether the agent will exit on an error while
+	// restoring the persistent cache
+	ExitOnErr bool
 }
 
 // New creates a new instance of Agent by parsing all the Kubernetes annotations.
@@ -330,7 +338,12 @@ func New(pod *corev1.Pod, patches []*jsonpatch.JsonPatchOperation) (*Agent, erro
 		return agent, err
 	}
 
-	agentCacheEnable, err := agent.agentCacheEnable()
+	agentCacheEnable, err := agent.cacheEnable()
+	if err != nil {
+		return agent, err
+	}
+
+	agentCacheExitOnErr, err := agent.cacheExitOnErr()
 	if err != nil {
 		return agent, err
 	}
@@ -339,6 +352,8 @@ func New(pod *corev1.Pod, patches []*jsonpatch.JsonPatchOperation) (*Agent, erro
 		Enable:           agentCacheEnable,
 		ListenerPort:     pod.Annotations[AnnotationAgentCacheListenerPort],
 		UseAutoAuthToken: pod.Annotations[AnnotationAgentCacheUseAutoAuthToken],
+		ExitOnErr:        agentCacheExitOnErr,
+		Persist:          agent.cachePersist(agentCacheEnable),
 	}
 
 	return agent, nil
@@ -419,6 +434,14 @@ func (a *Agent) Patch() ([]byte, error) {
 			"/spec/volumes")...)
 	}
 
+	// Add persistent cache volume if configured
+	if a.VaultAgentCache.Persist {
+		a.Patches = append(a.Patches, addVolumes(
+			a.Pod.Spec.Volumes,
+			[]corev1.Volume{a.cacheVolume()},
+			"/spec/volumes")...)
+	}
+
 	//Add Volume Mounts
 	for i, container := range a.Pod.Spec.Containers {
 		a.Patches = append(a.Patches, addVolumeMounts(

diff --git a/agent-inject/agent/annotations.go b/agent-inject/agent/annotations.go
@@ -205,6 +205,10 @@ const (
 	// AnnotationAgentCacheListenerPort configures the port the agent cache should listen on
 	AnnotationAgentCacheListenerPort = "vault.hashicorp.com/agent-cache-listener-port"
 
+	// AnnotationAgentCacheExitOnErr configures whether the agent will exit on an
+	// error while restoring the persistent cache
+	AnnotationAgentCacheExitOnErr = "vault.hashicorp.com/agent-cache-exit-on-err"
+
 	// AnnotationAgentCopyVolumeMounts is the name of the container or init container
 	// in the Pod whose volume mounts should be copied onto the Vault Agent init and
 	// sidecar containers. Ignores any Kubernetes service account token mounts.
@@ -363,6 +367,10 @@ func Init(pod *corev1.Pod, cfg AgentConfig) error {
 		pod.ObjectMeta.Annotations[AnnotationAgentCacheUseAutoAuthToken] = DefaultAgentCacheUseAutoAuthToken
 	}
 
+	if _, ok := pod.ObjectMeta.Annotations[AnnotationAgentCacheExitOnErr]; !ok {
+		pod.ObjectMeta.Annotations[AnnotationAgentCacheExitOnErr] = strconv.FormatBool(DefaultAgentCacheExitOnErr)
+	}
+
 	return nil
 }
 
@@ -541,7 +549,7 @@ func (a *Agent) setSecurityContext() (bool, error) {
 	return strconv.ParseBool(raw)
 }
 
-func (a *Agent) agentCacheEnable() (bool, error) {
+func (a *Agent) cacheEnable() (bool, error) {
 	raw, ok := a.Annotations[AnnotationAgentCacheEnable]
 	if !ok {
 		return false, nil
@@ -550,6 +558,22 @@ func (a *Agent) agentCacheEnable() (bool, error) {
 	return strconv.ParseBool(raw)
 }
 
+func (a *Agent) cachePersist(cacheEnabled bool) bool {
+	if cacheEnabled && a.PrePopulate && !a.PrePopulateOnly {
+		return true
+	}
+	return false
+}
+
+func (a *Agent) cacheExitOnErr() (bool, error) {
+	raw, ok := a.Annotations[AnnotationAgentCacheExitOnErr]
+	if !ok {
+		return false, nil
+	}
+
+	return strconv.ParseBool(raw)
+}
+
 func (a *Agent) authConfig() map[string]interface{} {
 	authConfig := make(map[string]interface{})
 

diff --git a/agent-inject/agent/config.go b/agent-inject/agent/config.go
@@ -85,7 +85,17 @@ type Listener struct {
 
 // Cache defines the configuration for the Vault Agent Cache
 type Cache struct {
-	UseAuthAuthToken string `json:"use_auto_auth_token"`
+	UseAutoAuthToken string        `json:"use_auto_auth_token"`
+	Persist          *CachePersist `json:"persist,omitempty"`
+}
+
+// CachePersist defines the configuration for persistent caching in Vault Agent
+type CachePersist struct {
+	Type                    string `json:"type"`
+	Path                    string `json:"path"`
+	KeepAfterImport         bool   `json:"keep_after_import,omitempty"`
+	ExitOnErr               bool   `json:"exit_on_err,omitempty"`
+	ServiceAccountTokenFile string `json:"service_account_token_file,omitempty"`
 }
 
 func (a *Agent) newTemplateConfigs() []*Template {
@@ -145,16 +155,27 @@ func (a *Agent) newConfig(init bool) ([]byte, error) {
 		Templates: a.newTemplateConfigs(),
 	}
 
-	if a.VaultAgentCache.Enable && !init {
-		config.Listener = []*Listener{
-			{
-				Type:       "tcp",
-				Address:    fmt.Sprintf("127.0.0.1:%s", a.VaultAgentCache.ListenerPort),
-				TLSDisable: true,
+	cacheListener := []*Listener{
+		{
+			Type:       "tcp",
+			Address:    fmt.Sprintf("127.0.0.1:%s", a.VaultAgentCache.ListenerPort),
+			TLSDisable: true,
+		},
+	}
+	if a.VaultAgentCache.Persist {
+		config.Listener = cacheListener
+		config.Cache = &Cache{
+			UseAutoAuthToken: a.VaultAgentCache.UseAutoAuthToken,
+			Persist: &CachePersist{
+				Type:      "kubernetes",
+				Path:      cacheVolumePath,
+				ExitOnErr: a.VaultAgentCache.ExitOnErr,
 			},
 		}
+	} else if a.VaultAgentCache.Enable && !a.PrePopulateOnly && !init {
+		config.Listener = cacheListener
 		config.Cache = &Cache{
-			UseAuthAuthToken: a.VaultAgentCache.UseAutoAuthToken,
+			UseAutoAuthToken: a.VaultAgentCache.UseAutoAuthToken,
 		}
 	}
 

diff --git a/agent-inject/agent/config_test.go b/agent-inject/agent/config_test.go
@@ -7,6 +7,8 @@ import (
 	"testing"
 
 	"github.com/mattbaird/jsonpatch"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestNewConfig(t *testing.T) {
@@ -306,8 +308,8 @@ func TestConfigVaultAgentCache(t *testing.T) {
 		t.Error("agent Cache should be enabled")
 	}
 
-	if config.Cache.UseAuthAuthToken != "force" {
-		t.Errorf("agent Cache use_auto_auth_token should be 'force', got %s instead", config.Cache.UseAuthAuthToken)
+	if config.Cache.UseAutoAuthToken != "force" {
+		t.Errorf("agent Cache use_auto_auth_token should be 'force', got %s instead", config.Cache.UseAutoAuthToken)
 	}
 
 	if config.Listener[0].Type != "tcp" {
@@ -322,3 +324,129 @@ func TestConfigVaultAgentCache(t *testing.T) {
 		t.Error("agent Cache listener TLS should be disabled")
 	}
 }
+
+func TestConfigVaultAgentCache_persistent(t *testing.T) {
+	tests := []struct {
+		name              string
+		annotations       map[string]string
+		expectedInitCache bool
+		expectedCache     *Cache
+		expectedListeners []*Listener
+	}{
+		{
+			name: "cache defaults",
+			annotations: map[string]string{
+				AnnotationAgentCacheEnable: "true",
+			},
+			expectedInitCache: true,
+			expectedCache: &Cache{
+				UseAutoAuthToken: "true",
+				Persist: &CachePersist{
+					Type: "kubernetes",
+					Path: "/vault/agent-cache",
+				},
+			},
+			expectedListeners: []*Listener{
+				{
+					Type:       "tcp",
+					Address:    "127.0.0.1:8200",
+					TLSDisable: true,
+				},
+			},
+		},
+		{
+			name: "exit on err",
+			annotations: map[string]string{
+				AnnotationAgentCacheEnable:    "true",
+				AnnotationAgentCacheExitOnErr: "true",
+			},
+			expectedInitCache: true,
+			expectedCache: &Cache{
+				UseAutoAuthToken: "true",
+				Persist: &CachePersist{
+					Type:      "kubernetes",
+					Path:      "/vault/agent-cache",
+					ExitOnErr: true,
+				},
+			},
+			expectedListeners: []*Listener{
+				{
+					Type:       "tcp",
+					Address:    "127.0.0.1:8200",
+					TLSDisable: true,
+				},
+			},
+		},
+		{
+			name: "just memory cache when only sidecar",
+			annotations: map[string]string{
+				AnnotationAgentCacheEnable: "true",
+				AnnotationAgentPrePopulate: "false",
+			},
+			expectedInitCache: false,
+			expectedCache: &Cache{
+				UseAutoAuthToken: "true",
+			},
+			expectedListeners: []*Listener{
+				{
+					Type:       "tcp",
+					Address:    "127.0.0.1:8200",
+					TLSDisable: true,
+				},
+			},
+		},
+		{
+			name: "no cache at all with only init container",
+			annotations: map[string]string{
+				AnnotationAgentCacheEnable:     "true",
+				AnnotationAgentPrePopulateOnly: "true",
+			},
+			expectedInitCache: false,
+			expectedCache:     nil,
+			expectedListeners: nil,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pod := testPod(tt.annotations)
+			var patches []*jsonpatch.JsonPatchOperation
+
+			agentConfig := AgentConfig{
+				"foobar-image", "http://foobar:8200", DefaultVaultAuthType, "test", "test", true, "100", "1000",
+				DefaultAgentRunAsSameUser, DefaultAgentSetSecurityContext, "",
+			}
+			err := Init(pod, agentConfig)
+			require.NoError(t, err, "got error initialising pod: %s", err)
+
+			agent, err := New(pod, patches)
+			require.NoError(t, err, "got error creating agent: %s", err)
+
+			initCfg, err := agent.newConfig(true)
+			require.NoError(t, err, "got error creating Vault config: %s", err)
+
+			initConfig := &Config{}
+			err = json.Unmarshal(initCfg, initConfig)
+			require.NoError(t, err, "got error unmarshalling Vault init config: %s", err)
+
+			if tt.expectedInitCache {
+				assert.Equal(t, tt.expectedCache, initConfig.Cache)
+				assert.Equal(t, tt.expectedListeners, initConfig.Listener)
+			} else {
+				assert.Nil(t, initConfig.Cache)
+				assert.Nil(t, initConfig.Listener)
+			}
+
+			sidecarCfg, err := agent.newConfig(false)
+			require.NoError(t, err, "got error creating Vault sidecar config: %s", err)
+
+			sidecarConfig := &Config{}
+			err = json.Unmarshal(sidecarCfg, sidecarConfig)
+			require.NoError(t, err, "got error unmarshalling Vault sidecar config: %s", err)
+
+			assert.Equal(t, tt.expectedCache, sidecarConfig.Cache)
+			assert.Equal(t, tt.expectedListeners, sidecarConfig.Listener)
+		})
+	}
+
+}
diff --git a/agent-inject/agent/container_init_sidecar.go b/agent-inject/agent/container_init_sidecar.go
@@ -57,6 +57,10 @@ func (a *Agent) ContainerInitSidecar() (corev1.Container, error) {
 		})
 	}
 
+	if a.VaultAgentCache.Persist {
+		volumeMounts = append(volumeMounts, a.cacheVolumeMount())
+	}
+
 	envs, err := a.ContainerEnvVars(true)
 	if err != nil {
 		return corev1.Container{}, err

diff --git a/agent-inject/agent/container_sidecar.go b/agent-inject/agent/container_sidecar.go
@@ -69,6 +69,10 @@ func (a *Agent) ContainerSidecar() (corev1.Container, error) {
 		})
 	}
 
+	if a.VaultAgentCache.Persist {
+		volumeMounts = append(volumeMounts, a.cacheVolumeMount())
+	}
+
 	envs, err := a.ContainerEnvVars(false)
 	if err != nil {
 		return corev1.Container{}, err