Redact username from vm args (#423)

lucko · Jul 20, 2024 · c0ffb19 · c0ffb19
1 parent 487cd0c
commit c0ffb19
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java b/spark-common/src/main/java/me/lucko/spark/common/platform/PlatformStatisticsProvider.java
@@ -43,6 +43,7 @@
 import java.lang.management.MemoryUsage;
 import java.lang.management.RuntimeMXBean;
 import java.util.Map;
+import java.util.regex.Pattern;
 
 public class PlatformStatisticsProvider {
     private final SparkPlatform platform;
@@ -55,6 +56,8 @@ public SystemStatistics getSystemStatistics() {
         RuntimeMXBean runtimeBean = ManagementFactory.getRuntimeMXBean();
         OperatingSystemInfo osInfo = OperatingSystemInfo.poll();
 
+        String vmArgs = String.join(" ", runtimeBean.getInputArguments());
+
         SystemStatistics.Builder builder = SystemStatistics.newBuilder()
                 .setCpu(SystemStatistics.Cpu.newBuilder()
                         .setThreads(Runtime.getRuntime().availableProcessors())
@@ -99,7 +102,7 @@ public SystemStatistics getSystemStatistics() {
                         .setVendor(System.getProperty("java.vendor", "unknown"))
                         .setVersion(System.getProperty("java.version", "unknown"))
                         .setVendorVersion(System.getProperty("java.vendor.version", "unknown"))
-                        .setVmArgs(String.join(" ", runtimeBean.getInputArguments()))
+                        .setVmArgs(VmArgRedactor.replace(vmArgs))
                         .build()
                 )
                 .setJvm(SystemStatistics.Jvm.newBuilder()
@@ -222,4 +225,17 @@ public static SparkProtos.RollingAverageValues rollingAvgProto(DoubleAverageInfo
                 .build();
     }
 
+    static final class VmArgRedactor {
+        private static final Pattern WINDOWS_USERNAME = Pattern.compile("C:\\\\Users\\\\\\w+");
+        private static final Pattern MACOS_USERNAME = Pattern.compile("/Users/\\w+");
+        private static final Pattern LINUX_USERNAME = Pattern.compile("/home/\\w+");
+
+        static String replace(String input) {
+            input = WINDOWS_USERNAME.matcher(input).replaceAll("C:\\\\Users\\\\<redacted>");
+            input = MACOS_USERNAME.matcher(input).replaceAll("/Users/<redacted>");
+            input = LINUX_USERNAME.matcher(input).replaceAll("/home/<redacted>");
+            return input;
+        }
+    }
+
 }