do not set WI env vars when NSA is enabled

cdapio · Nov 8, 2023 · 0f44398 · 0f44398
1 parent 9bd05f5
commit 0f44398
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 3 deletions.
diff --git a/...app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/DistributedPreviewManager.java b/...app-fabric/src/main/java/io/cdap/cdap/internal/app/preview/DistributedPreviewManager.java
@@ -235,6 +235,8 @@ public void run() {
                     String.format("%s:%s", localhost,
                         cConf.getInt(Constants.ArtifactLocalizer.PORT))
                 ));
+            twillPreparer = ((SecureTwillPreparer) twillPreparer)
+                .withNamespacedWorkloadIdentity(PreviewRunnerTwillRunnable.class.getSimpleName());
           }
 
           String priorityClass = cConf.get(Constants.Preview.CONTAINER_PRIORITY_CLASS_NAME);

diff --git a/...-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/TaskWorkerServiceLauncher.java b/...-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/TaskWorkerServiceLauncher.java
@@ -213,6 +213,8 @@ public void run() {
                     String.format("%s:%s", localhost,
                         cConf.getInt(Constants.ArtifactLocalizer.PORT))
                 ));
+            twillPreparer = ((SecureTwillPreparer) twillPreparer)
+                .withNamespacedWorkloadIdentity(TaskWorkerTwillRunnable.class.getSimpleName());
           }
 
           String priorityClass = cConf.get(Constants.TaskWorker.CONTAINER_PRIORITY_CLASS_NAME);

diff --git a/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java b/cdap-kubernetes/src/main/java/io/cdap/cdap/k8s/runtime/KubeTwillPreparer.java
@@ -189,6 +189,7 @@ class KubeTwillPreparer implements DependentTwillPreparer, StatefulTwillPreparer
   private final String resourcePrefix;
   private final Map<String, String> extraLabels;
   private final Map<String, SecretDiskRunnable> secretDiskRunnables;
+  private final Set<String> withNamespaceWorkloadIdentityRunnables;
   private final Map<String, V1SecurityContext> containerSecurityContexts;
   private final Map<String, Set<String>> readonlyDisks;
   private final Map<String, Map<String, String>> runnableConfigs;
@@ -240,6 +241,7 @@ class KubeTwillPreparer implements DependentTwillPreparer, StatefulTwillPreparer
     this.dependentRunnableNames = new HashSet<>();
     this.serviceAccountName = null;
     this.secretDiskRunnables = new HashMap<>();
+    this.withNamespaceWorkloadIdentityRunnables = new HashSet<>();
     this.containerSecurityContexts = new HashMap<>();
     this.readonlyDisks = new HashMap<>();
     this.runnableConfigs = runnables.stream()
@@ -368,6 +370,12 @@ public SecureTwillPreparer withSecretDisk(String runnableName, SecretDisk... sec
     return this;
   }
 
+  @Override
+  public SecureTwillPreparer withNamespacedWorkloadIdentity(String runnableName) {
+    withNamespaceWorkloadIdentityRunnables.add(runnableName);
+    return this;
+  }
+
   @Override
   public SecureTwillPreparer withSecurityContext(String runnableName,
       SecurityContext securityContext) {
@@ -1285,9 +1293,8 @@ private List<V1Container> createContainers(Map<String, RuntimeSpecification> run
     environs.put(JAVA_OPTS_KEY, jvmOpts);
     // Add workload identity environment variable if applicable.
     if (workloadIdentityEnabled && WorkloadIdentityUtil.shouldMountWorkloadIdentity(
-        cdapInstallNamespace,
-        programRuntimeNamespace,
-        workloadIdentityServiceAccount)) {
+        cdapInstallNamespace, programRuntimeNamespace, workloadIdentityServiceAccount)
+        && !withNamespaceWorkloadIdentityRunnables.contains(runnableName)) {
       V1EnvVar workloadIdentityEnvVar = WorkloadIdentityUtil.generateWorkloadIdentityEnvVar();
       environs.put(workloadIdentityEnvVar.getName(), workloadIdentityEnvVar.getValue());
     }
@@ -1314,6 +1321,13 @@ private List<V1Container> createContainers(Map<String, RuntimeSpecification> run
           .filter(entry -> !entry.getKey().equals(GCE_METADATA_HOST_ENV_VAR))
           .collect(Collectors.toMap(Map.Entry::getKey,
               Map.Entry::getValue));
+      // Add workload identity environment variable in the dependent runnable if applicable.
+      if (workloadIdentityEnabled && WorkloadIdentityUtil.shouldMountWorkloadIdentity(
+          cdapInstallNamespace, programRuntimeNamespace, workloadIdentityServiceAccount)
+          && !withNamespaceWorkloadIdentityRunnables.contains(name)) {
+        V1EnvVar workloadIdentityEnvVar = WorkloadIdentityUtil.generateWorkloadIdentityEnvVar();
+        envs.put(workloadIdentityEnvVar.getName(), workloadIdentityEnvVar.getValue());
+      }
       mounts = addSecreteVolMountIfNeeded(spec, volumeMounts);
       containers.add(
           createContainer(name, podInfo.getContainerImage(), podInfo.getImagePullPolicy(), workDir,

diff --git a/cdap-master-spi/src/main/java/io/cdap/cdap/master/spi/twill/SecureTwillPreparer.java b/cdap-master-spi/src/main/java/io/cdap/cdap/master/spi/twill/SecureTwillPreparer.java
@@ -44,4 +44,13 @@ public interface SecureTwillPreparer extends TwillPreparer {
   SecureTwillPreparer withSecurityContext(String runnableName,
       SecurityContext securityContext);
 
+  /**
+   * Runs the given runnable with namespace workload identity,
+   * this feature removes the GOOGLE_APPLICATION_CREDENTIALS environment variable
+   * to enable namespaced service accounts.
+   *
+   * @param runnableName name of the {@link TwillRunnable}
+   * @return this {@link TwillPreparer}
+   */
+  SecureTwillPreparer withNamespacedWorkloadIdentity(String runnableName);
 }