Extend CRD to support configuring container.spec.lifecycle.preStop.command in injected container

crd.yaml

@@ -6,12 +6,19 @@ spec:
   group: vaultwebhook.uswitch.com
   versions:
     - name: v1alpha1
+      # Each version can be enabled/disabled by Served flag.
       served: true
+      # One and only one version must be marked as the storage version.
       storage: true
       schema:
         openAPIV3Schema:
+          type: object
+          description: |-
+            A MutatingAdmissionController that will add the vault-creds container to your pod 
+            for you when your pod is created (assuming that vault webhook is enabled on your namespace
           properties:
             spec:
+              type: object
               properties:
                 database:
                   type: string
@@ -20,7 +27,29 @@ spec:
                 outputPath:
                   type: string
                 outputFile:
+                  type: string              
+                serviceAccount:
                   type: string
+                container:
+                  description: Specification of the container that will be created as part of this binding.
+                  type: object
+                  properties:
+                    lifecycle:
+                      description: Specification of the lifecycle hooks of the container. https://kubernetes.io/docs/concepts/containers/container-lifecycle-hooks/
+                      type: object
+                      properties:
+                        preStop:
+                          description: This hook is called immediately before a container is terminated due to an API request or management event such as a liveness/startup probe failure, preemption, resource contention and others
+                          type: object
+                          properties:
+                            exec:
+                              description: Executes a specific command, inside the cgroups and namespaces of the Container.
+                              type: object
+                              properties:
+                                command:
+                                  type: array
+                                  items:
+                                    type: string
   names:
     kind: DatabaseCredentialBinding
     plural: databasecredentialbindings

go.mod

@@ -71,4 +71,4 @@ require (
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.2.3 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect
-)
+)

go.sum

@@ -274,4 +274,4 @@ sigs.k8s.io/structured-merge-diff/v4 v4.2.3 h1:PRbqxJClWWYMNV1dhaG4NsibJbArud9kF
 sigs.k8s.io/structured-merge-diff/v4 v4.2.3/go.mod h1:qjx8mGObPmV2aSZepjQjbmb2ihdVs8cGKBraizNC69E=
 sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc=
 sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo=
-sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8=
+sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8=

pkg/apis/vaultwebhook.uswitch.com/v1alpha1/types.go

@@ -1,6 +1,7 @@
 package v1alpha1
 
 import (
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -15,11 +16,12 @@ type DatabaseCredentialBinding struct {
 }
 
 type DatabaseCredentialBindingSpec struct {
-	Database       string `json:"database"`
-	Role           string `json:"role"`
-	OutputPath     string `json:"outputPath"`
-	OutputFile     string `json:"outputFile"`
-	ServiceAccount string `json:"serviceAccount"`
+	Database       string    `json:"database"`
+	Role           string    `json:"role"`
+	OutputPath     string    `json:"outputPath"`
+	OutputFile     string    `json:"outputFile"`
+	ServiceAccount string    `json:"serviceAccount"`
+	Container      Container `json:"container,omitempty"`
 }
 
 // +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object
@@ -30,3 +32,23 @@ type DatabaseCredentialBindingList struct {
 
 	Items []DatabaseCredentialBinding `json:"items"`
 }
+
+type Container struct {
+	Lifecycle corev1.Lifecycle `json:"lifecycle,omitempty"`
+}
+
+/*
+https://pkg.go.dev/k8s.io/api/core/v1#LifecycleHandler
+Check if Container.Lifecycle.PreStop is valid. This is to avoid mishandling incomplete inputs like the below:
+
+	{	"Lifecycle": {
+			"PostStart": null,
+			"PreStop": {
+			  "Exec": null,  # <----- Missing Command!!
+			  "HTTPGet": null,"TCPSocket": null}}}
+*/
+func (c Container) HasValidPreStop() bool {
+	return c.Lifecycle.PreStop != nil &&
+		c.Lifecycle.PreStop.Exec != nil &&
+		len(c.Lifecycle.PreStop.Exec.Command) > 0
+}

vault.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/uswitch/vault-webhook/pkg/apis/vaultwebhook.uswitch.com/v1alpha1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 )
@@ -24,6 +25,8 @@ func addVault(pod *corev1.Pod, namespace string, databases []database) (patch []
 	initContainers := []corev1.Container{}
 	for _, databaseInfo := range databases {
 
+		vaultContainerSpec := databaseInfo.vaultContainer
+
 		database := databaseInfo.database
 		role := databaseInfo.role
 		serviceAccount := pod.Spec.ServiceAccountName
@@ -104,6 +107,9 @@ func addVault(pod *corev1.Pod, namespace string, databases []database) (patch []
 
 		initContainer := vaultContainer
 
+		// Configure Lifecycle Hooks if spec exists
+		vaultContainer = addLifecycleHook(vaultContainer, vaultContainerSpec)
+
 		jobLikeOwnerReferencesKinds := map[string]bool{"Job": true, "Workflow": true}
 		if len(pod.ObjectMeta.OwnerReferences) != 0 {
 			ownerKind := pod.ObjectMeta.OwnerReferences[0].Kind
@@ -112,6 +118,7 @@ func addVault(pod *corev1.Pod, namespace string, databases []database) (patch []
 			}
 		}
 
+		// Append the new Vault container spec into the Pod Spec generated by the client Deployment/Daemonset/etc
 		pod.Spec.Containers = append(pod.Spec.Containers, vaultContainer)
 
 		initContainer.Args = append(initContainer.Args, "--init")
@@ -197,3 +204,19 @@ func appendVolumeMountIfMissing(slice []corev1.VolumeMount, v corev1.VolumeMount
 	}
 	return append(slice, v)
 }
+
+// Conditionally set Lifecycle if it exists in containerSpec
+func addLifecycleHook(container corev1.Container, containerSpec v1alpha1.Container) corev1.Container {
+
+	// Check DatabaseCredentialBindingSpec.Container.Lifecycle is not empty
+	emptyLifecycle := corev1.Lifecycle{}
+	if containerSpec.Lifecycle != emptyLifecycle {
+
+		// Check for a complete PreStop hook
+		if containerSpec.HasValidPreStop() {
+			container.Lifecycle = &containerSpec.Lifecycle
+		}
+
+	}
+	return container
+}

vault_test.go

@@ -1,10 +1,12 @@
 package main
 
 import (
+	"fmt"
 	"strings"
 	"testing"
 
-	"k8s.io/api/core/v1"
+	"github.com/uswitch/vault-webhook/pkg/apis/vaultwebhook.uswitch.com/v1alpha1"
+	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
@@ -183,3 +185,68 @@ func TestVaultJobMode(t *testing.T) {
 		})
 	}
 }
+
+// Can we add a preStop hook to the vault container?
+func TestAddLifecyclePreStopHook(t *testing.T) {
+
+	// Define test cases
+	var tests = []struct {
+		scenario     string
+		lifecycleObj v1alpha1.Container
+		answer       bool
+	}{
+		{
+			scenario: "Test passing a complete lifecyle config",
+			lifecycleObj: v1alpha1.Container{
+				Lifecycle: v1.Lifecycle{
+					PreStop: &v1.LifecycleHandler{
+						Exec: &v1.ExecAction{
+							Command: []string{"echo", "hello"},
+						},
+					},
+				},
+			},
+			answer: true,
+		},
+		{
+			scenario: "Test passing an incomplete lifecycle config",
+			lifecycleObj: v1alpha1.Container{
+				Lifecycle: v1.Lifecycle{
+					PreStop: &v1.LifecycleHandler{
+						Exec: nil,
+					},
+				},
+			},
+			answer: false,
+		},
+		{
+			// v1alpha1.Container{}, comes from corev1.Container{} and this ALWAYS have a c.Lifecycle object. The latter, always has pointers to PostStart and PreStop handlers ( but no further down the struct since they are pointers )
+			// if our dcb input does not specify a container object, the received input will look like this: {Lifecycle:{PostStart:nil PreStop:nil}}
+			scenario: "Test passing no lifecycle config",
+			lifecycleObj: v1alpha1.Container{
+				Lifecycle: v1.Lifecycle{
+					PreStop: nil,
+				},
+			},
+			answer: false,
+		},
+	}
+
+	// Run tests
+	for _, tt := range tests {
+		// t.Run enables running "subtests", one for each table entry. These are shown separately when executing `go test -v`.
+		vaultContainer := v1.Container{} // Define a Vault sidecar Container
+		testname := fmt.Sprintf("%v", tt.scenario)
+		t.Run(testname, func(t *testing.T) {
+			ans := addLifecycleHook(vaultContainer, tt.lifecycleObj)
+
+			//log.Printf("%+v", ans)
+			isValid := ans.Lifecycle != nil && ans.Lifecycle.PreStop != nil && ans.Lifecycle.PreStop.Exec != nil && len(ans.Lifecycle.PreStop.Exec.Command) > 0
+
+			if isValid != tt.answer {
+				t.Errorf("got %v, want %v", isValid, tt.answer)
+			}
+		})
+	}
+
+}

webhook.go

@@ -38,10 +38,11 @@ type patchOperation struct {
 }
 
 type database struct {
-	database   string
-	role       string
-	outputPath string
-	outputFile string
+	database       string
+	role           string
+	outputPath     string
+	outputFile     string
+	vaultContainer v1alpha1.Container
 }
 
 func (srv webHookServer) serve(w http.ResponseWriter, r *http.Request) {
@@ -99,6 +100,7 @@ func (srv webHookServer) serve(w http.ResponseWriter, r *http.Request) {
 
 }
 
+// This handles the admission review sent by k8s and mutates the pod
 func (srv webHookServer) mutate(ar *v1beta1.AdmissionReview) *v1beta1.AdmissionResponse {
 	req := ar.Request
 
@@ -121,7 +123,9 @@ func (srv webHookServer) mutate(ar *v1beta1.AdmissionReview) *v1beta1.AdmissionR
 	log.Infof("AdmissionReview for Kind=%v, Namespace=%v Name=%v UID=%v patchOperation=%v UserInfo=%v",
 		ownerKind, req.Namespace, ownerName, req.UID, req.Operation, req.UserInfo)
 
+	// A list of ALL the bindings.
 	binds, err := srv.bindings.List()
+	log.Infof("[mutate] List of all bindings: %+v", binds)
 	if err != nil {
 		return &v1beta1.AdmissionResponse{
 			Result: &metav1.Status{
@@ -130,6 +134,7 @@ func (srv webHookServer) mutate(ar *v1beta1.AdmissionReview) *v1beta1.AdmissionR
 		}
 	}
 
+	// Filter out the bindings that are not in the target namespace
 	filteredBindings := filterBindings(binds, req.Namespace)
 	if len(filteredBindings) == 0 {
 		log.Infof("Skipping mutation for %s/%s, no database credential bindings in namespace", req.Namespace, ownerName)
@@ -138,6 +143,7 @@ func (srv webHookServer) mutate(ar *v1beta1.AdmissionReview) *v1beta1.AdmissionR
 		}
 	}
 
+	// Identify bindings with ServiceAccount field matching the pod's ServiceAccountName
 	databases := matchBindings(filteredBindings, pod.Spec.ServiceAccountName)
 	if len(databases) == 0 {
 		log.Infof("Skipping mutation for %s/%s due to policy check", req.Namespace, ownerName)
@@ -166,6 +172,7 @@ func (srv webHookServer) mutate(ar *v1beta1.AdmissionReview) *v1beta1.AdmissionR
 	}
 }
 
+// For all the bindings, we need to find the ones in the target namespace
 func filterBindings(bindings []v1alpha1.DatabaseCredentialBinding, namespace string) []v1alpha1.DatabaseCredentialBinding {
 	filteredBindings := []v1alpha1.DatabaseCredentialBinding{}
 	for _, binding := range bindings {
@@ -176,6 +183,12 @@ func filterBindings(bindings []v1alpha1.DatabaseCredentialBinding, namespace str
 	return filteredBindings
 }
 
+/*
+	    For all the bindings in the namespace, check which one has a ServiceeAccount that matches the pod's ServiceAccount
+		  - We could have multiple database specifications to be attached to a single pod.
+		  - This means that we could also have different VaultContainer specs for each DatabaseCredentialBinding.
+		  - As a consequence, to keep things consistent and easy to follow, we are appending into the `database` slice.
+*/
 func matchBindings(bindings []v1alpha1.DatabaseCredentialBinding, serviceAccount string) []database {
 	matchedBindings := []database{}
 	for _, binding := range bindings {
@@ -184,15 +197,27 @@ func matchBindings(bindings []v1alpha1.DatabaseCredentialBinding, serviceAccount
 			if output == "" {
 				output = "/etc/database"
 			}
-			matchedBindings = appendIfMissing(matchedBindings, database{role: binding.Spec.Role, database: binding.Spec.Database, outputPath: output, outputFile: binding.Spec.OutputFile})
+			log.Infof("[matchBindings] Printing content of Container: %+v", binding.Spec.Container)
+
+			matchedBindings = appendIfMissing(matchedBindings, database{
+				role:           binding.Spec.Role,
+				database:       binding.Spec.Database,
+				outputPath:     output,
+				outputFile:     binding.Spec.OutputFile,
+				vaultContainer: binding.Spec.Container,
+			})
 		}
 	}
 	return matchedBindings
 }
 
 func appendIfMissing(slice []database, d database) []database {
 	for _, ele := range slice {
-		if ele == d {
+		// No need to compare Container fields.
+		if ele.role == d.role &&
+			ele.database == d.database &&
+			ele.outputPath == d.outputPath &&
+			ele.outputFile == d.outputFile {
 			return slice
 		}
 	}