[api] Masks sensitive env vars in debug print out (#2657)

deepjavalibrary · Jun 15, 2023 · 8596e2e · 8596e2e
1 parent 3bcc6dc
commit 8596e2e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/api/src/main/java/ai/djl/engine/Engine.java b/api/src/main/java/ai/djl/engine/Engine.java
@@ -39,6 +39,7 @@
 import java.util.ServiceLoader;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.regex.Pattern;
 
 /**
  * The {@code Engine} interface is the base of the provided implementation for DJL.
@@ -59,6 +60,8 @@ public abstract class Engine {
     private static final Map<String, EngineProvider> ALL_ENGINES = new ConcurrentHashMap<>();
 
     private static final String DEFAULT_ENGINE = initEngine();
+    private static final Pattern PATTERN =
+            Pattern.compile("KEY|TOKEN|PASSWORD", Pattern.CASE_INSENSITIVE);
 
     private Device defaultDevice;
 
@@ -372,11 +375,11 @@ public String toString() {
     @SuppressWarnings("PMD.SystemPrintln")
     public static void debugEnvironment() {
         System.out.println("----------- System Properties -----------");
-        System.getProperties().forEach((k, v) -> System.out.println(k + ": " + v));
+        System.getProperties().forEach((k, v) -> print((String) k, v));
 
         System.out.println();
         System.out.println("--------- Environment Variables ---------");
-        Utils.getenv().forEach((k, v) -> System.out.println(k + ": " + v));
+        Utils.getenv().forEach(Engine::print);
 
         System.out.println();
         System.out.println("-------------- Directories --------------");
@@ -422,4 +425,11 @@ public static void debugEnvironment() {
             System.out.println(provider.getEngineName() + ": " + provider.getEngineRank());
         }
     }
+
+    private static void print(String key, Object value) {
+        if (PATTERN.matcher(key).find()) {
+            value = "*********";
+        }
+        System.out.println(key + ": " + value); // NOPMD
+    }
 }
diff --git a/api/src/test/java/ai/djl/DeviceTest.java b/api/src/test/java/ai/djl/DeviceTest.java
@@ -35,6 +35,7 @@ public void testDevice() {
         Device dev = Device.of("myDevice", 1);
         Assert.assertEquals(dev.getDeviceType(), "myDevice");
 
+        System.setProperty("test_key", "test");
         Engine.debugEnvironment();
     }