[secure-mode] make allowlist interface

deepjavalibrary · Jun 28, 2024 · d019e01 · d019e01
1 parent 4f53236
commit d019e01
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeAllowList.java b/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeAllowList.java
@@ -14,13 +14,18 @@
 
 import java.util.Set;
 
-/** A class for properties that are allowlisted in Secure Mode. */
-final class SecureModeAllowList {
+/** Properties that are explicitly allowlisted in Secure Mode. */
+interface SecureModeAllowList {
 
     public static final Set<String> PROPERTIES_ALLOWLIST =
             Set.of(
                     "engine",
                     "job_queue_size",
+                    "max_idle_time",
+                    "batch_size",
+                    "max_batch_delay",
+                    "minWorkers",
+                    "maxWorkers",
                     "option.entryPoint",
                     "option.task",
                     "option.model_id",
@@ -61,6 +66,4 @@ final class SecureModeAllowList {
                     "option.max_lora_rank",
                     "option.lora_extra_vocab_size",
                     "option.max_cpu_loras");
-
-    private SecureModeAllowList() {}
 }
diff --git a/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeUtils.java b/plugins/secure-mode/src/main/java/ai/djl/serving/plugins/securemode/SecureModeUtils.java
@@ -108,9 +108,9 @@ public static void validateSecurity(ModelInfo<?, ?> modelInfo) throws IOExceptio
      */
     private static void validateProperties(ModelInfo<?, ?> modelInfo, Set<String> allowedKeys) {
         Properties prop = modelInfo.getProperties();
-        Set<String> allowedKeysSet = new HashSet<>(allowedKeys);
+        allowedKeys = new HashSet<>(allowedKeys);
         for (String key : prop.stringPropertyNames()) {
-            if (!allowedKeysSet.contains(key)) {
+            if (!allowedKeys.contains(key)) {
                 throw new IllegalConfigurationException(
                         "Property " + key + " is prohibited from being set in Secure Mode.");
             }