Merge pull request #87865 from tedyu/container-visitor2

Allow container visitor to operate on selected container types

pkg/api/pod/util.go

@@ -26,28 +26,55 @@ import (
 	"k8s.io/kubernetes/pkg/security/apparmor"
 )
 
+// ContainerType signifies container type
+type ContainerType int
+
+// DefaultContainers defines default behavior: Iterate containers based on feature gates
+const DefaultContainers ContainerType = 0
+
+const (
+	// Containers is for normal containers
+	Containers ContainerType = 1 << iota
+	// InitContainers is for init containers
+	InitContainers
+	// EphemeralContainers is for ephemeral containers
+	EphemeralContainers
+)
+
+// AllContainers specifies that all containers be visited
+const AllContainers ContainerType = (InitContainers | Containers | EphemeralContainers)
+
 // ContainerVisitor is called with each container spec, and returns true
 // if visiting should continue.
-type ContainerVisitor func(container *api.Container) (shouldContinue bool)
+type ContainerVisitor func(container *api.Container, containerType ContainerType) (shouldContinue bool)
 
 // VisitContainers invokes the visitor function with a pointer to the container
 // spec of every container in the given pod spec. If visitor returns false,
 // visiting is short-circuited. VisitContainers returns true if visiting completes,
 // false if visiting was short-circuited.
-func VisitContainers(podSpec *api.PodSpec, visitor ContainerVisitor) bool {
+//
-	for i := range podSpec.InitContainers {
+// With the default mask (zero value or DefaultContainers) VisitContainers will visit all containers
-		if !visitor(&podSpec.InitContainers[i]) {
+// enabled by current feature gates. If mask is non-zero, VisitContainers will unconditionally visit
-			return false
+// container types specified by mask, and no feature gate checks will be performed.
+func VisitContainers(podSpec *api.PodSpec, mask ContainerType, visitor ContainerVisitor) bool {
+	if mask == DefaultContainers || (mask&InitContainers) > 0 {
+		for i := range podSpec.InitContainers {
+			if !visitor(&podSpec.InitContainers[i], InitContainers) {
+				return false
+			}
 		}
 	}
-	for i := range podSpec.Containers {
+	if mask == DefaultContainers || (mask&Containers) > 0 {
-		if !visitor(&podSpec.Containers[i]) {
+		for i := range podSpec.Containers {
-			return false
+			if !visitor(&podSpec.Containers[i], Containers) {
+				return false
+			}
 		}
 	}
-	if utilfeature.DefaultFeatureGate.Enabled(features.EphemeralContainers) {
+	if (mask == DefaultContainers && utilfeature.DefaultFeatureGate.Enabled(features.EphemeralContainers)) ||
+		(mask&EphemeralContainers) > 0 {
 		for i := range podSpec.EphemeralContainers {
-			if !visitor((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon)) {
+			if !visitor((*api.Container)(&podSpec.EphemeralContainers[i].EphemeralContainerCommon), EphemeralContainers) {
 				return false
 			}
 		}
@@ -68,7 +95,7 @@ func VisitPodSecretNames(pod *api.Pod, visitor Visitor) bool {
 			return false
 		}
 	}
-	VisitContainers(&pod.Spec, func(c *api.Container) bool {
+	VisitContainers(&pod.Spec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		return visitContainerSecretNames(c, visitor)
 	})
 	var source *api.VolumeSource
@@ -151,7 +178,7 @@ func visitContainerSecretNames(container *api.Container, visitor Visitor) bool {
 // Transitive references (e.g. pod -> pvc -> pv -> secret) are not visited.
 // Returns true if visiting completed, false if visiting was short-circuited.
 func VisitPodConfigmapNames(pod *api.Pod, visitor Visitor) bool {
-	VisitContainers(&pod.Spec, func(c *api.Container) bool {
+	VisitContainers(&pod.Spec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		return visitContainerConfigmapNames(c, visitor)
 	})
 	var source *api.VolumeSource
@@ -356,7 +383,7 @@ func dropDisabledFields(
 
 	if !utilfeature.DefaultFeatureGate.Enabled(features.VolumeSubpath) && !subpathInUse(oldPodSpec) {
 		// drop subpath from the pod if the feature is disabled and the old spec did not specify subpaths
-		VisitContainers(podSpec, func(c *api.Container) bool {
+		VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 			for i := range c.VolumeMounts {
 				c.VolumeMounts[i].SubPath = ""
 			}
@@ -369,7 +396,7 @@ func dropDisabledFields(
 
 	if (!utilfeature.DefaultFeatureGate.Enabled(features.VolumeSubpath) || !utilfeature.DefaultFeatureGate.Enabled(features.VolumeSubpathEnvExpansion)) && !subpathExprInUse(oldPodSpec) {
 		// drop subpath env expansion from the pod if either of the subpath features is disabled and the old spec did not specify subpath env expansion
-		VisitContainers(podSpec, func(c *api.Container) bool {
+		VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 			for i := range c.VolumeMounts {
 				c.VolumeMounts[i].SubPathExpr = ""
 			}
@@ -379,7 +406,7 @@ func dropDisabledFields(
 
 	if !utilfeature.DefaultFeatureGate.Enabled(features.StartupProbe) && !startupProbeInUse(oldPodSpec) {
 		// drop startupProbe from all containers if the feature is disabled
-		VisitContainers(podSpec, func(c *api.Container) bool {
+		VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 			c.StartupProbe = nil
 			return true
 		})
@@ -421,7 +448,7 @@ func dropDisabledRunAsGroupField(podSpec, oldPodSpec *api.PodSpec) {
 		if podSpec.SecurityContext != nil {
 			podSpec.SecurityContext.RunAsGroup = nil
 		}
-		VisitContainers(podSpec, func(c *api.Container) bool {
+		VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 			if c.SecurityContext != nil {
 				c.SecurityContext.RunAsGroup = nil
 			}
@@ -435,7 +462,7 @@ func dropDisabledRunAsGroupField(podSpec, oldPodSpec *api.PodSpec) {
 func dropDisabledProcMountField(podSpec, oldPodSpec *api.PodSpec) {
 	if !utilfeature.DefaultFeatureGate.Enabled(features.ProcMountType) && !procMountInUse(oldPodSpec) {
 		defaultProcMount := api.DefaultProcMount
-		VisitContainers(podSpec, func(c *api.Container) bool {
+		VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 			if c.SecurityContext != nil && c.SecurityContext.ProcMount != nil {
 				// The ProcMount field was improperly forced to non-nil in 1.12.
 				// If the feature is disabled, and the existing object is not using any non-default values, and the ProcMount field is present in the incoming object, force to the default value.
@@ -471,7 +498,7 @@ func subpathInUse(podSpec *api.PodSpec) bool {
 	}
 
 	var inUse bool
-	VisitContainers(podSpec, func(c *api.Container) bool {
+	VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		for i := range c.VolumeMounts {
 			if len(c.VolumeMounts[i].SubPath) > 0 {
 				inUse = true
@@ -521,7 +548,7 @@ func procMountInUse(podSpec *api.PodSpec) bool {
 	}
 
 	var inUse bool
-	VisitContainers(podSpec, func(c *api.Container) bool {
+	VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		if c.SecurityContext == nil || c.SecurityContext.ProcMount == nil {
 			return true
 		}
@@ -609,7 +636,7 @@ func runAsGroupInUse(podSpec *api.PodSpec) bool {
 	}
 
 	var inUse bool
-	VisitContainers(podSpec, func(c *api.Container) bool {
+	VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		if c.SecurityContext != nil && c.SecurityContext.RunAsGroup != nil {
 			inUse = true
 			return false
@@ -627,7 +654,7 @@ func subpathExprInUse(podSpec *api.PodSpec) bool {
 	}
 
 	var inUse bool
-	VisitContainers(podSpec, func(c *api.Container) bool {
+	VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		for i := range c.VolumeMounts {
 			if len(c.VolumeMounts[i].SubPathExpr) > 0 {
 				inUse = true
@@ -647,7 +674,7 @@ func startupProbeInUse(podSpec *api.PodSpec) bool {
 	}
 
 	var inUse bool
-	VisitContainers(podSpec, func(c *api.Container) bool {
+	VisitContainers(podSpec, AllContainers, func(c *api.Container, containerType ContainerType) bool {
 		if c.StartupProbe != nil {
 			inUse = true
 			return false

pkg/api/pod/util_test.go

@@ -41,11 +41,13 @@ func TestVisitContainers(t *testing.T) {
 		description string
 		haveSpec    *api.PodSpec
 		wantNames   []string
+		mask        ContainerType
 	}{
 		{
 			"empty podspec",
 			&api.PodSpec{},
 			[]string{},
+			DefaultContainers,
 		},
 		{
 			"regular containers",
@@ -56,6 +58,7 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"c1", "c2"},
+			DefaultContainers,
 		},
 		{
 			"init containers",
@@ -66,6 +69,7 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"i1", "i2"},
+			DefaultContainers,
 		},
 		{
 			"regular and init containers",
@@ -80,6 +84,7 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"i1", "i2", "c1", "c2"},
+			DefaultContainers,
 		},
 		{
 			"ephemeral containers",
@@ -93,6 +98,7 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"c1", "c2", "e1"},
+			DefaultContainers,
 		},
 		{
 			"all container types",
@@ -111,6 +117,26 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"i1", "i2", "c1", "c2", "e1", "e2"},
+			DefaultContainers,
+		},
+		{
+			"all container types with init and regular container types chosen",
+			&api.PodSpec{
+				Containers: []api.Container{
+					{Name: "c1"},
+					{Name: "c2"},
+				},
+				InitContainers: []api.Container{
+					{Name: "i1"},
+					{Name: "i2"},
+				},
+				EphemeralContainers: []api.EphemeralContainer{
+					{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e1"}},
+					{EphemeralContainerCommon: api.EphemeralContainerCommon{Name: "e2"}},
+				},
+			},
+			[]string{"i1", "i2", "c1", "c2"},
+			Containers | InitContainers,
 		},
 		{
 			"dropping fields",
@@ -129,12 +155,13 @@ func TestVisitContainers(t *testing.T) {
 				},
 			},
 			[]string{"i1", "i2", "c1", "c2", "e1", "e2"},
+			DefaultContainers,
 		},
 	}
 
 	for _, tc := range testCases {
 		gotNames := []string{}
-		VisitContainers(tc.haveSpec, func(c *api.Container) bool {
+		VisitContainers(tc.haveSpec, tc.mask, func(c *api.Container, containerType ContainerType) bool {
 			gotNames = append(gotNames, c.Name)
 			if c.SecurityContext != nil {
 				c.SecurityContext = nil

pkg/registry/core/pod/strategy.go

@@ -387,7 +387,7 @@ func LogLocation(
 
 func podHasContainerWithName(pod *api.Pod, containerName string) bool {
 	var hasContainer bool
-	podutil.VisitContainers(&pod.Spec, func(c *api.Container) bool {
+	podutil.VisitContainers(&pod.Spec, podutil.DefaultContainers, func(c *api.Container, containerType podutil.ContainerType) bool {
 		if c.Name == containerName {
 			hasContainer = true
 			return false
@@ -554,7 +554,7 @@ func validateContainer(container string, pod *api.Pod) (string, error) {
 			return "", errors.NewBadRequest(fmt.Sprintf("a container name must be specified for pod %s", pod.Name))
 		default:
 			var containerNames []string
-			podutil.VisitContainers(&pod.Spec, func(c *api.Container) bool {
+			podutil.VisitContainers(&pod.Spec, podutil.DefaultContainers, func(c *api.Container, containerType podutil.ContainerType) bool {
 				containerNames = append(containerNames, c.Name)
 				return true
 			})

pkg/security/podsecuritypolicy/provider.go

@@ -111,7 +111,7 @@ func (s *simpleProvider) MutatePod(pod *api.Pod) error {
 	}
 
 	var retErr error
-	podutil.VisitContainers(&pod.Spec, func(c *api.Container) bool {
+	podutil.VisitContainers(&pod.Spec, podutil.AllContainers, func(c *api.Container, containerType podutil.ContainerType) bool {
 		retErr = s.mutateContainer(pod, c)
 		if retErr != nil {
 			return false