feat: allow proxies to be injected using a native sidecar

docs/book/src/topics/service-account-labels-and-annotations.md

@@ -21,9 +21,9 @@ All annotations are optional. If the annotation is not specified, the default va
 | `azure.workload.identity/service-account-token-expiration` | **(Takes precedence if the service account is also annotated)** Represents the `expirationSeconds` field for the projected service account token. It is an optional field that the user might want to configure this to prevent any downtime caused by errors during service account token refresh. Kubernetes service account token expiry will not be correlated with AAD tokens. AAD tokens will expire in 24 hours after they are issued. | `3600` (acceptable range: `3600 - 86400`) |
 | `azure.workload.identity/skip-containers`                  | Represents a semi-colon-separated list of containers (e.g. `container1;container2`) to skip adding projected service account token volume. By default, the projected service account token volume will be added to all containers.                                                                                                                                                                                                            |                                           |
 | `azure.workload.identity/inject-proxy-sidecar`             | Injects a proxy init container and proxy sidecar into the pod. The proxy sidecar is used to intercept token requests to IMDS and acquire an AAD token on behalf of the user with federated identity credential.                                                                                                                                                                                                                               | `true`                                    |
+| `azure.workload.identity/use-native-sidecar`               | If injected then the proxy sidecar should be injected as a native sidecar init container instead of a container. Native sidecars are terminated when if pod exits, for instance when running as a cron-job.                                                                                                                                                                                                                                   | `false`                                   |
 | `azure.workload.identity/proxy-sidecar-port`               | Represents the port of the proxy sidecar.                                                                                                                                                                                                                                                                                                                                                                                                     | `8000`                                    |
 
-
 ## Service Account
 
 ### Annotations

pkg/webhook/consts.go

@@ -17,6 +17,8 @@ const (
 	SkipContainersAnnotation = "azure.workload.identity/skip-containers"
 	// InjectProxySidecarAnnotation represents the annotation to be used to inject proxy sidecar into the pod
 	InjectProxySidecarAnnotation = "azure.workload.identity/inject-proxy-sidecar"
+	// UseNativeSidecarAnnotation represents the annotation to be used to inject the proxy as a native sidecar into the pod
+	UseNativeSidecarAnnotation = "azure.workload.identity/use-native-sidecar"
 	// ProxySidecarPortAnnotation represents the annotation to be used to specify the port for proxy sidecar
 	ProxySidecarPortAnnotation = "azure.workload.identity/proxy-sidecar-port"
 

pkg/webhook/webhook.go

@@ -153,8 +153,14 @@ func (m *podMutator) Handle(ctx context.Context, req admission.Request) (respons
 			return admission.Errored(http.StatusBadRequest, err)
 		}
 
+		useNativeSidecar := useNativeProxySidecar(pod)
+
 		pod.Spec.InitContainers = m.injectProxyInitContainer(pod.Spec.InitContainers, proxyPort)
-		pod.Spec.Containers = m.injectProxySidecarContainer(pod.Spec.Containers, proxyPort)
+		if useNativeSidecar {
+			pod.Spec.InitContainers = m.injectProxySidecarContainer(pod.Spec.InitContainers, proxyPort, ptr.To(corev1.ContainerRestartPolicyAlways))
+		} else {
+			pod.Spec.Containers = m.injectProxySidecarContainer(pod.Spec.Containers, proxyPort, nil)
+		}
 	}
 
 	// get service account token expiration
@@ -227,13 +233,18 @@ func (m *podMutator) injectProxyInitContainer(containers []corev1.Container, pro
 	return containers
 }
 
-func (m *podMutator) injectProxySidecarContainer(containers []corev1.Container, proxyPort int32) []corev1.Container {
+func (m *podMutator) injectProxySidecarContainer(containers []corev1.Container, proxyPort int32, restartPolicy *corev1.ContainerRestartPolicy) []corev1.Container {
+	imageRepository := strings.Join([]string{ProxyImageRegistry, ProxySidecarImageName}, "/")
 	for _, container := range containers {
 		if container.Name == ProxySidecarContainerName {
-			return containers
+			// proxy-init starts with proxy, so we append ":" to imageRepository when checking
+			if strings.HasPrefix(container.Image, fmt.Sprintf("%s:", imageRepository)) || container.Name == ProxySidecarContainerName {
+				return containers
+			}
 		}
 	}
 	logLevel := currentLogLevel() // run the proxy at the same log level as the webhook
+
 	containers = append([]corev1.Container{{
 		Name:            ProxySidecarContainerName,
 		Image:           m.proxyImage,
@@ -266,6 +277,7 @@ func (m *podMutator) injectProxySidecarContainer(containers []corev1.Container,
 			ReadOnlyRootFilesystem: ptr.To(true),
 			RunAsNonRoot:           ptr.To(true),
 		},
+		RestartPolicy: restartPolicy,
 	}}, containers...)
 
 	return containers
@@ -279,6 +291,14 @@ func shouldInjectProxySidecar(pod *corev1.Pod) bool {
 	return ok
 }
 
+func useNativeProxySidecar(pod *corev1.Pod) bool {
+	if len(pod.Annotations) == 0 {
+		return false
+	}
+	_, ok := pod.Annotations[UseNativeSidecarAnnotation]
+	return ok
+}
+
 // getSkipContainers gets the list of containers to skip based on the annotation
 func getSkipContainers(pod *corev1.Pod) map[string]struct{} {
 	skipContainers := pod.Annotations[SkipContainersAnnotation]

pkg/webhook/webhook_test.go

@@ -1150,7 +1150,7 @@ func TestInjectProxySidecarContainer(t *testing.T) {
 	m := &podMutator{proxyImage: imageURL}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			containers := m.injectProxySidecarContainer(test.containers, proxyPort)
+			containers := m.injectProxySidecarContainer(test.containers, proxyPort, nil)
 			if !reflect.DeepEqual(containers, test.expectedContainers) {
 				t.Errorf("expected: %v, got: %v", test.expectedContainers, containers)
 			}
@@ -1197,6 +1197,45 @@ func TestShouldInjectProxySidecar(t *testing.T) {
 	}
 }
 
+func TestShouldUseNativeSidecar(t *testing.T) {
+	tests := []struct {
+		name     string
+		pod      *corev1.Pod
+		expected bool
+	}{
+		{
+			name: "pod not annotated",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "pod",
+				},
+			},
+			expected: false,
+		},
+		{
+			name: "pod is annotated with azure.workload.identity/use-native-sidecar=true",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "pod",
+					Annotations: map[string]string{
+						UseNativeSidecarAnnotation: "true",
+					},
+				},
+			},
+			expected: true,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			actual := useNativeProxySidecar(test.pod)
+			if actual != test.expected {
+				t.Fatalf("expected: %v, got: %v", test.expected, actual)
+			}
+		})
+	}
+}
+
 func TestGetProxyPort(t *testing.T) {
 	type args struct {
 		pod *corev1.Pod