Add support to mitigate log4j vulnerabilities from user imported log4… (

#626) * Add support to mitigate log4j vulnerabilities from user imported log4j libraries Signed-off-by: Frank Gu <[email protected]> * update test case Signed-off-by: Frank Gu <[email protected]> * update java agent syntax Signed-off-by: Frank Gu <[email protected]> * Release TonY v0.4.12 Signed-off-by: Frank Gu <[email protected]>
tony-framework · Dec 15, 2021 · d198bed · d198bed
1 parent 05bc8c9
commit d198bed
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 2 deletions.
diff --git a/build.gradle b/build.gradle
@@ -92,7 +92,7 @@ ext.deps = [
 
 allprojects {
   group = "com.linkedin.tony"
-  project.version = "0.4.11"
+  project.version = "0.4.12"
 }
 
 task sourcesJar(type: Jar) {

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java b/tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java
@@ -155,6 +155,8 @@ private TonyConfigurationKeys() {
   public static final String TASK_EXECUTOR_JVM_OPTS = TONY_TASK_PREFIX + "executor.jvm.opts";
   public static final String DEFAULT_TASK_EXECUTOR_JVM_OPTS = "-Xmx1536m";
 
+  public static final String TASK_EXECUTOR_JAVA_AGENT = TONY_TASK_PREFIX + "executor.java.agent";
+
   public static final String TASK_HEARTBEAT_INTERVAL_MS = TONY_TASK_PREFIX + "heartbeat-interval-ms";
   public static final int DEFAULT_TASK_HEARTBEAT_INTERVAL_MS = 1000;
 

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonySession.java b/tony-core/src/main/java/com/linkedin/tony/TonySession.java
@@ -429,14 +429,24 @@ public TonySession build() {
     }
 
     public Builder setTaskExecutorJVMArgs(String jvmArgs) {
-      this.jvmArgs = jvmArgs;
+      this.jvmArgs = appendDefaultJVMArgs(jvmArgs);
       return this;
     }
 
     public Builder setTonyConf(Configuration tonyConf) {
       this.tonyConf = tonyConf;
       return this;
     }
+
+    // Appends default jvm arguments to task executor job
+    private String appendDefaultJVMArgs(String jvmArgs) {
+      jvmArgs += " -Dlog4j2.formatMsgNoLookups=true";
+      String tonyTaskExecutorJavaAgent = tonyConf.get(TonyConfigurationKeys.TASK_EXECUTOR_JAVA_AGENT, "");
+      if (!tonyTaskExecutorJavaAgent.isEmpty()) {
+        jvmArgs += " -javaagent:" + tonyTaskExecutorJavaAgent;
+      }
+      return jvmArgs;
+    }
   }
 
   @VisibleForTesting

diff --git a/tony-core/src/test/java/com/linkedin/tony/TestTonyConfigurationFields.java b/tony-core/src/test/java/com/linkedin/tony/TestTonyConfigurationFields.java
@@ -49,6 +49,7 @@ public void initializeMemberVariables() {
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_TRAINING_STAGE);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_HADOOP_LOCATION);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.APPLICATION_HADOOP_CLASSPATH);
+    configurationPropsToSkipCompare.add(TonyConfigurationKeys.TASK_EXECUTOR_JAVA_AGENT);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.TENSORBOARD_LOG_DIR);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.PYTHON_EXEC_PATH);
     configurationPropsToSkipCompare.add(TonyConfigurationKeys.TB_VCORE);