From 30369c2ad5f04424eba0e62c765ff8d19d648877 Mon Sep 17 00:00:00 2001
From: James Powis <james.powis@charter.com>
Date: Fri, 5 Oct 2018 15:49:19 -0600
Subject: [PATCH] Drop volumeMounts from initContainers if SAToken

Signed-off-by: James Powis <powisj@gmail.com>
---
 pkg/restore/pod_action.go      | 31 +++++++++++++++++++++++++
 pkg/restore/pod_action_test.go | 41 ++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/pkg/restore/pod_action.go b/pkg/restore/pod_action.go
index f18231408..2ec930bba 100644
--- a/pkg/restore/pod_action.go
+++ b/pkg/restore/pod_action.go
@@ -112,5 +112,36 @@ func (a *podAction) Execute(obj runtime.Unstructured, restore *api.Restore) (run
 		return nil, nil, err
 	}
 
+	a.logger.Debug("iterating over init containers")
+	err = collections.ForEach(spec, "initContainers", func(container map[string]interface{}) error {
+		var newVolumeMounts []interface{}
+		err := collections.ForEach(container, "volumeMounts", func(volumeMount map[string]interface{}) error {
+			name, err := collections.GetString(volumeMount, "name")
+			if err != nil {
+				return err
+			}
+
+			a.logger.WithField("volumeMount", name).Debug("Checking volumeMount")
+			if strings.HasPrefix(name, serviceAccountName+"-token-") {
+				a.logger.WithField("volumeMount", name).Debug("Excluding volumeMount")
+			} else {
+				a.logger.WithField("volumeMount", name).Debug("Preserving volumeMount")
+				newVolumeMounts = append(newVolumeMounts, volumeMount)
+			}
+
+			return nil
+		})
+		if err != nil {
+			return err
+		}
+
+		container["volumeMounts"] = newVolumeMounts
+
+		return nil
+	})
+	if err != nil {
+		return nil, nil, err
+	}
+
 	return obj, nil, nil
 }
diff --git a/pkg/restore/pod_action_test.go b/pkg/restore/pod_action_test.go
index b3b1dde9f..7dfd9e7b8 100644
--- a/pkg/restore/pod_action_test.go
+++ b/pkg/restore/pod_action_test.go
@@ -43,18 +43,21 @@ func TestPodActionExecute(t *testing.T) {
 				WithSpec("serviceAccountName", "foo").
 				WithSpecField("volumes", []interface{}{}).
 				WithSpecField("containers", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{}).
 				Unstructured,
 			expectedErr: false,
 			expectedRes: NewTestUnstructured().WithName("pod-1").WithSpec("foo").
 				WithSpec("serviceAccountName", "foo").
 				WithSpecField("volumes", []interface{}{}).
 				WithSpecField("containers", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{}).
 				Unstructured,
 		},
 		{
 			name: "volumes matching prefix ServiceAccount-token- should be deleted",
 			obj: NewTestUnstructured().WithName("pod-1").
 				WithSpec("serviceAccountName", "foo").
+				WithSpecField("initContainers", []interface{}{}).
 				WithSpecField("volumes", []interface{}{
 					map[string]interface{}{"name": "foo"},
 					map[string]interface{}{"name": "foo-token-foo"},
@@ -62,6 +65,7 @@ func TestPodActionExecute(t *testing.T) {
 			expectedErr: false,
 			expectedRes: NewTestUnstructured().WithName("pod-1").
 				WithSpec("serviceAccountName", "foo").
+				WithSpecField("initContainers", []interface{}{}).
 				WithSpecField("volumes", []interface{}{
 					map[string]interface{}{"name": "foo"},
 				}).WithSpecField("containers", []interface{}{}).Unstructured,
@@ -71,6 +75,7 @@ func TestPodActionExecute(t *testing.T) {
 			obj: NewTestUnstructured().WithName("svc-1").
 				WithSpec("serviceAccountName", "foo").
 				WithSpecField("volumes", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{}).
 				WithSpecField("containers", []interface{}{
 					map[string]interface{}{
 						"volumeMounts": []interface{}{
@@ -88,6 +93,7 @@ func TestPodActionExecute(t *testing.T) {
 			expectedRes: NewTestUnstructured().WithName("svc-1").
 				WithSpec("serviceAccountName", "foo").
 				WithSpecField("volumes", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{}).
 				WithSpecField("containers", []interface{}{
 					map[string]interface{}{
 						"volumeMounts": []interface{}{
@@ -99,6 +105,41 @@ func TestPodActionExecute(t *testing.T) {
 				}).
 				Unstructured,
 		},
+		{
+			name: "initContainer volumeMounts matching prefix ServiceAccount-token- should be deleted",
+			obj: NewTestUnstructured().WithName("svc-1").
+				WithSpec("serviceAccountName", "foo").
+				WithSpecField("volumes", []interface{}{}).
+				WithSpecField("containers", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{
+					map[string]interface{}{
+						"volumeMounts": []interface{}{
+							map[string]interface{}{
+								"name": "foo",
+							},
+							map[string]interface{}{
+								"name": "foo-token-foo",
+							},
+						},
+					},
+				}).
+				Unstructured,
+			expectedErr: false,
+			expectedRes: NewTestUnstructured().WithName("svc-1").
+				WithSpec("serviceAccountName", "foo").
+				WithSpecField("volumes", []interface{}{}).
+				WithSpecField("containers", []interface{}{}).
+				WithSpecField("initContainers", []interface{}{
+					map[string]interface{}{
+						"volumeMounts": []interface{}{
+							map[string]interface{}{
+								"name": "foo",
+							},
+						},
+					},
+				}).
+				Unstructured,
+		},
 	}
 
 	for _, test := range tests {