Merge pull request #3 from shipt/feature/manual-rolearn-env

Gtoken arn from env var
shipt · Jan 6, 2022 · 934c038 · 934c038
2 parents b2ba37f + fbe89b0
commit 934c038
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 7 deletions.
diff --git a/cmd/gtoken-webhook/main.go b/cmd/gtoken-webhook/main.go
@@ -50,6 +50,7 @@ const (
 	awsWebIdentityTokenFile = "AWS_WEB_IDENTITY_TOKEN_FILE"
 	awsRoleArn              = "AWS_ROLE_ARN"
 	awsRoleSessionName      = "AWS_ROLE_SESSION_NAME"
+	awsGtokenEnvVarName     = "GTOKEN_AWS_ROLE_ARN"
 )
 
 var (
@@ -203,15 +204,35 @@ func (mw *mutatingWebhook) mutatePod(pod *corev1.Pod, ns string, dryRun bool) er
 		return nil
 	}
 
-	// get service account AWS Role ARN annotation
-	roleArn, ok, err := mw.getAwsRoleArn(pod.Spec.ServiceAccountName, ns)
-	if err != nil {
-		return err
+	manualRoleArn := ""
+	for _, container := range pod.Spec.Containers {
+		for _, env := range container.Env {
+			if env.Name == awsGtokenEnvVarName && env.Value != "" {
+				manualRoleArn = env.Value
+				break
+			}
+		}
+
+		if manualRoleArn != "" {
+			break
+		}
 	}
-	if !ok {
-		logger.Debug("skipping pods with Service Account without AWS Role ARN annotation")
-		return nil
+
+	var roleArn = manualRoleArn
+	if roleArn == "" { // check role arn on SA if not specified in ENV
+		var ok bool
+		var err error
+		// get service account AWS Role ARN annotation
+		roleArn, ok, err = mw.getAwsRoleArn(pod.Spec.ServiceAccountName, ns)
+		if err != nil {
+			return err
+		}
+		if !ok {
+			logger.Debug("skipping pods with Service Account without AWS Role ARN annotation")
+			return nil
+		}
 	}
+
 	// mutate Pod init containers
 	initContainersMutated := mw.mutateContainers(pod.Spec.InitContainers, roleArn)
 	if initContainersMutated {

diff --git a/cmd/gtoken-webhook/main_test.go b/cmd/gtoken-webhook/main_test.go
@@ -318,6 +318,118 @@ func Test_mutatingWebhook_mutatePod(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "manual role arn",
+			fields: fields{
+				image:      "doitintl/gtoken:test",
+				pullPolicy: "Always",
+				volumeName: "test-volume-name",
+				volumePath: "/test-volume-path",
+				tokenFile:  "test-token",
+			},
+			args: args{
+				pod: &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{annotationInjectKey: "true"},
+					},
+					Spec: corev1.PodSpec{
+						Containers: []corev1.Container{
+							{
+								Name:  "TestContainer",
+								Image: "test-image",
+								Env: []corev1.EnvVar{
+									{
+										Name:  "GTOKEN_AWS_ROLE_ARN",
+										Value: "arn:aws:iam::123456789012:role/testrole",
+									},
+								},
+							},
+						},
+						ServiceAccountName: "test-sa",
+					},
+				},
+				ns:                 "test-namespace",
+				serviceAccountName: "test-sa",
+				annotations:        map[string]string{},
+			},
+			wantedPod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Annotations: map[string]string{annotationInjectKey: "true"},
+				},
+				Spec: corev1.PodSpec{
+					InitContainers: []corev1.Container{
+						{
+							Name:    "generate-gcp-id-token",
+							Image:   "doitintl/gtoken:test",
+							Command: []string{"/gtoken", "--file=/test-volume-path/test-token", "--refresh=false"},
+							Resources: corev1.ResourceRequirements{
+								Requests: corev1.ResourceList{
+									corev1.ResourceCPU:    resource.MustParse(requestsCPU),
+									corev1.ResourceMemory: resource.MustParse(requestsMemory),
+								},
+								Limits: corev1.ResourceList{
+									corev1.ResourceCPU:    resource.MustParse(limitsCPU),
+									corev1.ResourceMemory: resource.MustParse(limitsMemory),
+								},
+							},
+							VolumeMounts: []corev1.VolumeMount{
+								{
+									Name:      "test-volume-name",
+									MountPath: "/test-volume-path",
+								},
+							},
+							ImagePullPolicy: "Always",
+						},
+					},
+					Containers: []corev1.Container{
+						{
+							Name:         "TestContainer",
+							Image:        "test-image",
+							VolumeMounts: []corev1.VolumeMount{{Name: "test-volume-name", MountPath: "/test-volume-path"}},
+							Env: []corev1.EnvVar{
+								{Name: awsGtokenEnvVarName, Value: "arn:aws:iam::123456789012:role/testrole"},
+								{Name: awsWebIdentityTokenFile, Value: "/test-volume-path/test-token"},
+								{Name: awsRoleArn, Value: "arn:aws:iam::123456789012:role/testrole"},
+								{Name: awsRoleSessionName, Value: "gtoken-webhook-" + strings.Repeat("0", 16)},
+							},
+						},
+						{
+							Name:    "update-gcp-id-token",
+							Image:   "doitintl/gtoken:test",
+							Command: []string{"/gtoken", "--file=/test-volume-path/test-token", "--refresh=true"},
+							Resources: corev1.ResourceRequirements{
+								Requests: corev1.ResourceList{
+									corev1.ResourceCPU:    resource.MustParse(requestsCPU),
+									corev1.ResourceMemory: resource.MustParse(requestsMemory),
+								},
+								Limits: corev1.ResourceList{
+									corev1.ResourceCPU:    resource.MustParse(limitsCPU),
+									corev1.ResourceMemory: resource.MustParse(limitsMemory),
+								},
+							},
+							VolumeMounts: []corev1.VolumeMount{
+								{
+									Name:      "test-volume-name",
+									MountPath: "/test-volume-path",
+								},
+							},
+							ImagePullPolicy: "Always",
+						},
+					},
+					Volumes: []corev1.Volume{
+						{
+							Name: "test-volume-name",
+							VolumeSource: corev1.VolumeSource{
+								EmptyDir: &corev1.EmptyDirVolumeSource{
+									Medium: corev1.StorageMediumMemory,
+								},
+							},
+						},
+					},
+					ServiceAccountName: "test-sa",
+				},
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {