From d019e014ff911f6c592704e074f39bc162653281 Mon Sep 17 00:00:00 2001
From: Ethan Zhang <ez3141@gmail.com>
Date: Fri, 28 Jun 2024 20:57:29 +0000
Subject: [PATCH] [secure-mode] make allowlist interface

---
 .../plugins/securemode/SecureModeAllowList.java       | 11 +++++++----
 .../serving/plugins/securemode/SecureModeUtils.java   |  4 ++--
 2 files changed, 9 insertions(+), 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
index ebdbc7334..54ef8d6a9 100644
--- 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
index 8d59c190b..9da88e61e 100644
--- 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.");
             }