Support Revoking Vault Token on Shutdown (hashicorp#67)

* Write Vault Token to the usual location

As documented by the default token helper: https://www.vaultproject.io/docs/commands/token-helper/

* Add annotation

* Add Lifecycle to container

* Fix gofmt

* Add option to automatically revoke tokens for all pods

* Add flags to revoke command

* Make flags only inside if clause

agent-inject/agent/agent.go

@@ -52,9 +52,17 @@ type Agent struct {
 	PrePopulate bool
 
 	// PrePopulateOnly controls whether an init container is the _only_ container
-	//added to the request.
+	// added to the request.
 	PrePopulateOnly bool
 
+	// RevokeOnShutdown controls whether a sidecar container will attempt to revoke its Vault
+	// token on shutting down.
+	RevokeOnShutdown bool
+
+	// RevokeGrace controls after receiving the signal for pod
+	// termination that the container will attempt to revoke its own Vault token.
+	RevokeGrace uint64
+
 	// RequestsCPU is the requested minimum CPU amount required  when being scheduled to deploy.
 	RequestsCPU string
 
@@ -205,6 +213,16 @@ func New(pod *corev1.Pod, patches []*jsonpatch.JsonPatchOperation) (*Agent, erro
 		return agent, err
 	}
 
+	agent.RevokeOnShutdown, err = agent.revokeOnShutdown()
+	if err != nil {
+		return agent, err
+	}
+
+	agent.RevokeGrace, err = agent.revokeGrace()
+	if err != nil {
+		return agent, err
+	}
+
 	agent.Vault.TLSSkipVerify, err = agent.tlsSkipVerify()
 	if err != nil {
 		return agent, err
@@ -364,3 +382,23 @@ func serviceaccount(pod *corev1.Pod) (string, string) {
 	}
 	return serviceAccountName, serviceAccountPath
 }
+
+func (a *Agent) vaultCliFlags() []string {
+	flags := []string{
+		fmt.Sprintf("-address=%s", a.Vault.Address),
+	}
+
+	if a.Vault.CACert != "" {
+		flags = append(flags, fmt.Sprintf("-ca-cert=%s", a.Vault.CACert))
+	}
+
+	if a.Vault.ClientCert != "" {
+		flags = append(flags, fmt.Sprintf("-client-cert=%s", a.Vault.ClientCert))
+	}
+
+	if a.Vault.ClientKey != "" {
+		flags = append(flags, fmt.Sprintf("-client-key=%s", a.Vault.ClientKey))
+	}
+
+	return flags
+}

agent-inject/agent/annotations.go

@@ -79,6 +79,15 @@ const (
 	// AnnotationAgentRequestsMem sets the requested memory amount on the Vault Agent containers.
 	AnnotationAgentRequestsMem = "vault.hashicorp.com/agent-requests-mem"
 
+	// AnnotationAgentRevokeOnShutdown controls whether a sidecar container will revoke its
+	// own Vault token before shutting down. If you are using a custom agent template, you must
+	// make sure it's written to `/home/vault/.vault-token`. Only supported for sidecar containers.
+	AnnotationAgentRevokeOnShutdown = "vault.hashicorp.com/agent-revoke-on-shutdown"
+
+	// AnnotationAgentRevokeGrace sets the number of seconds after receiving the signal for pod
+	// termination that the container will attempt to revoke its own Vault token. Defaults to 5s.
+	AnnotationAgentRevokeGrace = "vault.hashicorp.com/agent-revoke-grace"
+
 	// AnnotationVaultNamespace is the Vault namespace where secrets can be found.
 	AnnotationVaultNamespace = "vault.hashicorp.com/namespace"
 
@@ -134,7 +143,7 @@ const (
 // Init configures the expected annotations required to create a new instance
 // of Agent.  This should be run before running new to ensure all annotations are
 // present.
-func Init(pod *corev1.Pod, image, address, authPath, namespace string) error {
+func Init(pod *corev1.Pod, image, address, authPath, namespace string, revokeOnShutdown bool) error {
 	if pod == nil {
 		return errors.New("pod is empty")
 	}
@@ -190,6 +199,14 @@ func Init(pod *corev1.Pod, image, address, authPath, namespace string) error {
 		pod.ObjectMeta.Annotations[AnnotationAgentRequestsMem] = DefaultResourceRequestMem
 	}
 
+	if _, ok := pod.ObjectMeta.Annotations[AnnotationAgentRevokeOnShutdown]; !ok {
+		pod.ObjectMeta.Annotations[AnnotationAgentRevokeOnShutdown] = strconv.FormatBool(revokeOnShutdown)
+	}
+
+	if _, ok := pod.ObjectMeta.Annotations[AnnotationAgentRevokeGrace]; !ok {
+		pod.ObjectMeta.Annotations[AnnotationAgentRevokeGrace] = strconv.Itoa(DefaultRevokeGrace)
+	}
+
 	if _, ok := pod.ObjectMeta.Annotations[AnnotationVaultLogLevel]; !ok {
 		pod.ObjectMeta.Annotations[AnnotationVaultLogLevel] = DefaultAgentLogLevel
 	}
@@ -267,6 +284,24 @@ func (a *Agent) prePopulateOnly() (bool, error) {
 	return strconv.ParseBool(raw)
 }
 
+func (a *Agent) revokeOnShutdown() (bool, error) {
+	raw, ok := a.Annotations[AnnotationAgentRevokeOnShutdown]
+	if !ok {
+		return false, nil
+	}
+
+	return strconv.ParseBool(raw)
+}
+
+func (a *Agent) revokeGrace() (uint64, error) {
+	raw, ok := a.Annotations[AnnotationAgentRevokeGrace]
+	if !ok {
+		return 0, nil
+	}
+
+	return strconv.ParseUint(raw, 10, 64)
+}
+
 func (a *Agent) tlsSkipVerify() (bool, error) {
 	raw, ok := a.Annotations[AnnotationVaultTLSSkipVerify]
 	if !ok {

agent-inject/agent/annotations_test.go

@@ -14,7 +14,7 @@ func TestInitCanSet(t *testing.T) {
 	annotations := make(map[string]string)
 	pod := testPod(annotations)
 
-	err := Init(pod, "foobar-image", "http://foobar:8200", "test", "test")
+	err := Init(pod, "foobar-image", "http://foobar:8200", "test", "test", true)
 	if err != nil {
 		t.Errorf("got error, shouldn't have: %s", err)
 	}
@@ -26,6 +26,7 @@ func TestInitCanSet(t *testing.T) {
 		{annotationKey: AnnotationVaultService, annotationValue: "http://foobar:8200"},
 		{annotationKey: AnnotationAgentImage, annotationValue: "foobar-image"},
 		{annotationKey: AnnotationAgentRequestNamespace, annotationValue: "test"},
+		{annotationKey: AnnotationAgentRevokeOnShutdown, annotationValue: "true"},
 	}
 
 	for _, tt := range tests {
@@ -45,7 +46,7 @@ func TestInitDefaults(t *testing.T) {
 	annotations := make(map[string]string)
 	pod := testPod(annotations)
 
-	err := Init(pod, "", "http://foobar:8200", "test", "test")
+	err := Init(pod, "", "http://foobar:8200", "test", "test", true)
 	if err != nil {
 		t.Errorf("got error, shouldn't have: %s", err)
 	}
@@ -74,7 +75,7 @@ func TestInitError(t *testing.T) {
 	annotations := make(map[string]string)
 	pod := testPod(annotations)
 
-	err := Init(pod, "image", "", "authPath", "namespace")
+	err := Init(pod, "image", "", "authPath", "namespace", true)
 	if err == nil {
 		t.Error("expected error no address, got none")
 	}
@@ -84,7 +85,7 @@ func TestInitError(t *testing.T) {
 		t.Errorf("expected '%s' error, got %s", errMsg, err)
 	}
 
-	err = Init(pod, "image", "address", "", "namespace")
+	err = Init(pod, "image", "address", "", "namespace", true)
 	if err == nil {
 		t.Error("expected error no authPath, got none")
 	}
@@ -94,7 +95,7 @@ func TestInitError(t *testing.T) {
 		t.Errorf("expected '%s' error, got %s", errMsg, err)
 	}
 
-	err = Init(pod, "image", "address", "authPath", "")
+	err = Init(pod, "image", "address", "authPath", "", true)
 	if err == nil {
 		t.Error("expected error for no namespace, got none")
 	}
@@ -375,6 +376,24 @@ func TestCouldErrorAnnotations(t *testing.T) {
 		{AnnotationVaultTLSSkipVerify, "tRuE", false},
 		{AnnotationVaultTLSSkipVerify, "fAlSe", false},
 		{AnnotationVaultTLSSkipVerify, "", false},
+
+		{AnnotationAgentRevokeOnShutdown, "true", true},
+		{AnnotationAgentRevokeOnShutdown, "false", true},
+		{AnnotationAgentRevokeOnShutdown, "TRUE", true},
+		{AnnotationAgentRevokeOnShutdown, "FALSE", true},
+		{AnnotationAgentRevokeOnShutdown, "0", true},
+		{AnnotationAgentRevokeOnShutdown, "1", true},
+		{AnnotationAgentRevokeOnShutdown, "t", true},
+		{AnnotationAgentRevokeOnShutdown, "f", true},
+		{AnnotationAgentRevokeOnShutdown, "tRuE", false},
+		{AnnotationAgentRevokeOnShutdown, "fAlSe", false},
+		{AnnotationAgentRevokeOnShutdown, "", false},
+
+		{AnnotationAgentRevokeGrace, "5", true},
+		{AnnotationAgentRevokeGrace, "0", true},
+		{AnnotationAgentRevokeGrace, "01", true},
+		{AnnotationAgentRevokeGrace, "-1", false},
+		{AnnotationAgentRevokeGrace, "foobar", false},
 	}
 
 	for i, tt := range tests {
@@ -394,7 +413,7 @@ func TestCouldErrorAnnotations(t *testing.T) {
 func TestInitEmptyPod(t *testing.T) {
 	var pod *corev1.Pod
 
-	err := Init(pod, "foobar-image", "http://foobar:8200", "test", "test")
+	err := Init(pod, "foobar-image", "http://foobar:8200", "test", "test", true)
 	if err == nil {
 		t.Errorf("got no error, shouldn have")
 	}

agent-inject/agent/container_sidecar.go

@@ -2,6 +2,8 @@ package agent
 
 import (
 	"fmt"
+	"strings"
+
 	"k8s.io/apimachinery/pkg/api/resource"
 
 	"github.com/hashicorp/vault/sdk/helper/pointerutil"
@@ -15,6 +17,7 @@ const (
 	DefaultResourceRequestCPU = "250m"
 	DefaultResourceRequestMem = "64Mi"
 	DefaultContainerArg       = "echo ${VAULT_CONFIG?} | base64 -d > /tmp/config.json && vault agent -config=/tmp/config.json"
+	DefaultRevokeGrace        = 5
 	DefaultAgentLogLevel      = "info"
 )
 
@@ -63,6 +66,8 @@ func (a *Agent) ContainerSidecar() (corev1.Container, error) {
 		return corev1.Container{}, err
 	}
 
+	lifecycle := a.createLifecycle()
+
 	return corev1.Container{
 		Name:      "vault-agent",
 		Image:     a.ImageName,
@@ -73,6 +78,7 @@ func (a *Agent) ContainerSidecar() (corev1.Container, error) {
 			RunAsGroup:   pointerutil.Int64Ptr(1000),
 			RunAsNonRoot: pointerutil.BoolPtr(true),
 		},
+		Lifecycle:    &lifecycle,
 		VolumeMounts: volumeMounts,
 		Command:      []string{"/bin/sh", "-ec"},
 		Args:         []string{arg},
@@ -125,3 +131,21 @@ func parseQuantity(raw string) (resource.Quantity, error) {
 
 	return resource.ParseQuantity(raw)
 }
+
+// This should only be run for a sidecar container
+func (a *Agent) createLifecycle() corev1.Lifecycle {
+	lifecycle := corev1.Lifecycle{}
+
+	if a.RevokeOnShutdown {
+		flags := a.vaultCliFlags()
+		flags = append(flags, "-self")
+
+		lifecycle.PreStop = &corev1.Handler{
+			Exec: &corev1.ExecAction{
+				Command: []string{"/bin/sh", "-c", fmt.Sprintf("/bin/sleep %d && /bin/vault token revoke %s", a.RevokeGrace, strings.Join(flags[:], " "))},
+			},
+		}
+	}
+
+	return lifecycle
+}

agent-inject/agent/container_sidecar_test.go

@@ -15,7 +15,7 @@ func TestContainerSidecar(t *testing.T) {
 	pod := testPod(annotations)
 	var patches []*jsonpatch.JsonPatchOperation
 
-	err := Init(pod, "foobar-image", "http://foobar:1234", "test", "test")
+	err := Init(pod, "foobar-image", "http://foobar:1234", "test", "test", false)
 	if err != nil {
 		t.Errorf("got error, shouldn't have: %s", err)
 	}
@@ -76,6 +76,66 @@ func TestContainerSidecar(t *testing.T) {
 	}
 }
 
+func TestContainerSidecarRevokeHook(t *testing.T) {
+	trueString := "true"
+	falseString := "false"
+
+	tests := []struct {
+		revokeFlag       bool
+		revokeAnnotation *string
+		expectedPresence bool
+	}{
+		{revokeFlag: true, revokeAnnotation: nil, expectedPresence: true},
+		{revokeFlag: false, revokeAnnotation: nil, expectedPresence: false},
+		{revokeFlag: true, revokeAnnotation: &trueString, expectedPresence: true},
+		{revokeFlag: true, revokeAnnotation: &falseString, expectedPresence: false},
+		{revokeFlag: false, revokeAnnotation: &trueString, expectedPresence: true},
+		{revokeFlag: false, revokeAnnotation: &falseString, expectedPresence: false},
+	}
+
+	for _, tt := range tests {
+		t.Run("revoke test", func(t *testing.T) {
+			var revokeAnnotation string
+
+			annotations := map[string]string{
+				AnnotationVaultRole: "foobar",
+			}
+
+			if tt.revokeAnnotation == nil {
+				revokeAnnotation = "<absent>"
+			} else {
+				annotations[AnnotationAgentRevokeOnShutdown] = *tt.revokeAnnotation
+			}
+
+			pod := testPod(annotations)
+			var patches []*jsonpatch.JsonPatchOperation
+
+			err := Init(pod, "foobar-image", "http://foobar:1234", "test", "test", tt.revokeFlag)
+			if err != nil {
+				t.Errorf("got error, shouldn't have: %s", err)
+			}
+
+			agent, err := New(pod, patches)
+			if err := agent.Validate(); err != nil {
+				t.Errorf("agent validation failed, it shouldn't have: %s", err)
+			}
+
+			container, err := agent.ContainerSidecar()
+			if err != nil {
+				t.Errorf("creating container sidecar failed, it shouldn't have: %s", err)
+			}
+
+			if tt.expectedPresence && container.Lifecycle.PreStop == nil {
+				t.Errorf("revoke flag was %t and annotation was %s but preStop hook was absent when it was expected to be present", tt.revokeFlag, revokeAnnotation)
+			}
+
+			if !tt.expectedPresence && container.Lifecycle.PreStop != nil {
+				t.Errorf("revoke flag was %t and annotation was %s but preStop hook was present when it was expected to not be present", tt.revokeFlag, revokeAnnotation)
+			}
+		})
+	}
+}
+
 func TestContainerSidecarConfigMap(t *testing.T) {
 	// None of these custom configs should matter since
 	// we have AnnotationAgentConfigMap set
@@ -98,7 +158,7 @@ func TestContainerSidecarConfigMap(t *testing.T) {
 	pod := testPod(annotations)
 	var patches []*jsonpatch.JsonPatchOperation
 
-	err := Init(pod, "foobar-image", "http://foobar:1234", "test", "test")
+	err := Init(pod, "foobar-image", "http://foobar:1234", "test", "test", true)
 	if err != nil {
 		t.Errorf("got error, shouldn't have: %s", err)
 	}

agent-inject/handler.go

@@ -41,6 +41,7 @@ type Handler struct {
 	ImageVault        string
 	Clientset         *kubernetes.Clientset
 	Log               hclog.Logger
+	RevokeOnShutdown  bool
 }
 
 // Handle is the http.HandlerFunc implementation that actually handles the
@@ -135,7 +136,7 @@ func (h *Handler) Mutate(req *v1beta1.AdmissionRequest) *v1beta1.AdmissionRespon
 
 	h.Log.Debug("setting default annotations..")
 	var patches []*jsonpatch.JsonPatchOperation
-	err = agent.Init(&pod, h.ImageVault, h.VaultAddress, h.VaultAuthPath, req.Namespace)
+	err = agent.Init(&pod, h.ImageVault, h.VaultAddress, h.VaultAuthPath, req.Namespace, h.RevokeOnShutdown)
 	if err != nil {
 		err := fmt.Errorf("error adding default annotations: %s", err)
 		return admissionError(err)