[merger/security_context] Fix AddCapabilitiesToContainer (#562)

controllers/datadogagent/merger/security_context.go

@@ -12,7 +12,7 @@ import (
 
 // SecurityContextManager use to add Security Context settings to containers.
 type SecurityContextManager interface {
-	// Add capabilities to a container in the PodTemplate.
+	// AddCapabilitiesToContainer Adds capabilities to a container in the PodTemplate.
 	AddCapabilitiesToContainer(capabilities []corev1.Capability, containerName commonv1.AgentContainerName)
 }
 
@@ -28,17 +28,20 @@ type securityContextManagerImpl struct {
 }
 
 func (impl *securityContextManagerImpl) AddCapabilitiesToContainer(capabilities []corev1.Capability, containerName commonv1.AgentContainerName) {
-	for _, container := range impl.podTmpl.Spec.Containers {
+	for i, container := range impl.podTmpl.Spec.Containers {
 		if container.Name == string(containerName) {
 			if container.SecurityContext == nil {
-				container.SecurityContext = &corev1.SecurityContext{
+				impl.podTmpl.Spec.Containers[i].SecurityContext = &corev1.SecurityContext{
 					Capabilities: &corev1.Capabilities{
 						Add: capabilities,
 					},
 				}
 			} else {
 				// TODO add deduplication
-				container.SecurityContext.Capabilities.Add = append(container.SecurityContext.Capabilities.Add, capabilities...)
+				impl.podTmpl.Spec.Containers[i].SecurityContext.Capabilities.Add = append(
+					impl.podTmpl.Spec.Containers[i].SecurityContext.Capabilities.Add,
+					capabilities...,
+				)
 			}
 			return
 		}

controllers/datadogagent/merger/security_context_test.go

@@ -0,0 +1,121 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016-present Datadog, Inc.
+
+package merger
+
+import (
+	"testing"
+
+	commonv1 "github.com/DataDog/datadog-operator/apis/datadoghq/common/v1"
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+)
+
+func TestAddCapabilitiesToContainer(t *testing.T) {
+	tests := []struct {
+		name                   string
+		existingContainers     []corev1.Container
+		capabilities           []corev1.Capability
+		addToContainerWithName commonv1.AgentContainerName
+		expectedCapabilities   map[commonv1.AgentContainerName][]corev1.Capability
+	}{
+		{
+			name: "Add to container without capabilities defined",
+			existingContainers: []corev1.Container{
+				{
+					Name:            string(commonv1.TraceAgentContainerName),
+					SecurityContext: nil,
+				},
+			},
+			capabilities: []corev1.Capability{
+				"AUDIT_CONTROL",
+				"AUDIT_READ",
+			},
+			addToContainerWithName: commonv1.TraceAgentContainerName,
+			expectedCapabilities: map[commonv1.AgentContainerName][]corev1.Capability{
+				commonv1.TraceAgentContainerName: {
+					"AUDIT_CONTROL",
+					"AUDIT_READ",
+				},
+			},
+		},
+		{
+			name: "Add to container with some capabilities already defined",
+			existingContainers: []corev1.Container{
+				{
+					Name: string(commonv1.TraceAgentContainerName),
+					SecurityContext: &corev1.SecurityContext{
+						Capabilities: &corev1.Capabilities{
+							Add: []corev1.Capability{
+								"AUDIT_CONTROL",
+							},
+						},
+					},
+				},
+			},
+			capabilities: []corev1.Capability{
+				"AUDIT_READ",
+			},
+			addToContainerWithName: commonv1.TraceAgentContainerName,
+			expectedCapabilities: map[commonv1.AgentContainerName][]corev1.Capability{
+				commonv1.TraceAgentContainerName: {
+					"AUDIT_CONTROL",
+					"AUDIT_READ",
+				},
+			},
+		},
+		{
+			name: "Add to specific container when there are multiple defined",
+			existingContainers: []corev1.Container{
+				{
+					Name:            string(commonv1.TraceAgentContainerName),
+					SecurityContext: nil,
+				},
+				{
+					Name:            string(commonv1.SystemProbeContainerName),
+					SecurityContext: nil,
+				},
+			},
+			capabilities: []corev1.Capability{
+				"AUDIT_CONTROL",
+				"AUDIT_READ",
+			},
+			addToContainerWithName: commonv1.SystemProbeContainerName,
+			expectedCapabilities: map[commonv1.AgentContainerName][]corev1.Capability{
+				commonv1.TraceAgentContainerName: nil,
+				commonv1.SystemProbeContainerName: {
+					"AUDIT_CONTROL",
+					"AUDIT_READ",
+				},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			securityContextManager := securityContextManagerImpl{
+				podTmpl: &corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						Containers: test.existingContainers,
+					},
+				},
+			}
+
+			securityContextManager.AddCapabilitiesToContainer(test.capabilities, test.addToContainerWithName)
+
+			for _, container := range securityContextManager.podTmpl.Spec.Containers {
+				expectedCapabilities := test.expectedCapabilities[commonv1.AgentContainerName(container.Name)]
+				if len(expectedCapabilities) > 0 {
+					assert.Equal(t, expectedCapabilities, container.SecurityContext.Capabilities.Add)
+				} else {
+					assert.True(t, container.SecurityContext == nil ||
+						container.SecurityContext.Capabilities == nil ||
+						len(container.SecurityContext.Capabilities.Add) == 0)
+				}
+			}
+
+		})
+	}
+}