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]>

Author: UWFrankGu

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) {

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;
 

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

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);