refresh ml key

Signed-off-by: Jing Zhang <[email protected]>

Author: jngz-es

common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

@@ -350,4 +350,8 @@ private MLCommonsSettings() {}
     // Feature flag for enabling telemetry static metric collection job -- MLStatsJobProcessor
     public static final Setting<Boolean> ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED = Setting
         .boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Final);
+
+    // This setting is for refreshing ML key from the ml-config index to the memory.
+    public static final Setting<Boolean> ML_COMMONS_KEY_REFRESH_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.key_refresh_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
 }

common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java

@@ -10,6 +10,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_KEY_REFRESH_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_LOCAL_MODEL_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED;
@@ -50,6 +51,7 @@ public class MLFeatureEnabledSetting {
 
     private volatile Boolean isMetricCollectionEnabled;
     private volatile Boolean isStaticMetricCollectionEnabled;
+    private volatile Boolean isKeyRefreshEnabled;
 
     private final List<SettingsChangeListener> listeners = new ArrayList<>();
 
@@ -66,6 +68,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isRagSearchPipelineEnabled = ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED.get(settings);
         isMetricCollectionEnabled = ML_COMMONS_METRIC_COLLECTION_ENABLED.get(settings);
         isStaticMetricCollectionEnabled = ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED.get(settings);
+        isKeyRefreshEnabled = ML_COMMONS_KEY_REFRESH_ENABLED.get(settings);
 
         clusterService
             .getClusterSettings()
@@ -88,6 +91,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService
             .getClusterSettings()
             .addSettingsUpdateConsumer(MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED, it -> isRagSearchPipelineEnabled = it);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_KEY_REFRESH_ENABLED, it -> isKeyRefreshEnabled = it);
     }
 
     /**
@@ -178,6 +182,14 @@ public boolean isStaticMetricCollectionEnabled() {
         return isStaticMetricCollectionEnabled;
     }
 
+    /**
+     * Whether the key refresh feature is enabled.
+     * @return whether the key refresh is enabled.
+     */
+    public boolean isKeyRefreshEnabled() {
+        return isKeyRefreshEnabled;
+    }
+
     @VisibleForTesting
     public void notifyMultiTenancyListeners(boolean isEnabled) {
         for (SettingsChangeListener listener : listeners) {

ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java

@@ -121,6 +121,18 @@ public String generateMasterKey() {
         return Base64.getEncoder().encodeToString(keyBytes);
     }
 
+    // Refresh the key from the config index.
+    public void refreshMasterKey(ActionListener<Boolean> listener) {
+        try {
+            // Currently, we only handle no tenant case.
+            initMasterKeyFromIndex(null);
+            listener.onResponse(true);
+        } catch (Exception e) {
+            log.info("Refreshing ML key failed.");
+            listener.onFailure(e);
+        }
+    }
+
     private JceMasterKey createJceMasterKey(String tenantId) {
         byte[] bytes = Base64.getDecoder().decode(tenantMasterKeys.get(Objects.requireNonNullElse(tenantId, DEFAULT_TENANT_ID)));
         return JceMasterKey.getInstance(new SecretKeySpec(bytes, "AES"), "Custom", "", "AES/GCM/NOPADDING");
@@ -130,6 +142,10 @@ private void initMasterKey(String tenantId) {
         if (tenantMasterKeys.containsKey(Objects.requireNonNullElse(tenantId, DEFAULT_TENANT_ID))) {
             return;
         }
+        initMasterKeyFromIndex(tenantId);
+    }
+
+    private void initMasterKeyFromIndex(String tenantId) {
         String masterKeyId = MASTER_KEY;
         if (tenantId != null) {
             masterKeyId = MASTER_KEY + "_" + hashString(tenantId);

plugin/src/main/java/org/opensearch/ml/action/syncup/TransportSyncUpOnNodeAction.java

@@ -23,11 +23,13 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.ml.common.MLTask;
 import org.opensearch.ml.common.MLTaskState;
 import org.opensearch.ml.common.MLTaskType;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.sync.MLSyncUpAction;
 import org.opensearch.ml.common.transport.sync.MLSyncUpInput;
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodeRequest;
@@ -36,6 +38,7 @@
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodesResponse;
 import org.opensearch.ml.engine.MLEngine;
 import org.opensearch.ml.engine.ModelHelper;
+import org.opensearch.ml.engine.encryptor.EncryptorImpl;
 import org.opensearch.ml.engine.utils.FileUtils;
 import org.opensearch.ml.model.MLModelCacheHelper;
 import org.opensearch.ml.model.MLModelManager;
@@ -66,6 +69,8 @@ public class TransportSyncUpOnNodeAction extends
     private volatile Integer mlTaskTimeout;
 
     private final MLModelCacheHelper mlModelCacheHelper;
+    private final EncryptorImpl encryptor;
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     @Inject
     public TransportSyncUpOnNodeAction(
@@ -80,7 +85,9 @@ public TransportSyncUpOnNodeAction(
         Client client,
         NamedXContentRegistry xContentRegistry,
         MLEngine mlEngine,
-        MLModelCacheHelper mlModelCacheHelper
+        MLModelCacheHelper mlModelCacheHelper,
+        EncryptorImpl encryptor,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
     ) {
         super(
             MLSyncUpAction.NAME,
@@ -103,6 +110,8 @@ public TransportSyncUpOnNodeAction(
         this.xContentRegistry = xContentRegistry;
         this.mlEngine = mlEngine;
         this.mlModelCacheHelper = mlModelCacheHelper;
+        this.encryptor = encryptor;
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
 
         this.mlTaskTimeout = ML_COMMONS_ML_TASK_TIMEOUT_IN_SECONDS.get(settings);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_ML_TASK_TIMEOUT_IN_SECONDS, it -> { mlTaskTimeout = it; });
@@ -182,6 +191,14 @@ private MLSyncUpNodeResponse createSyncUpNodeResponse(MLSyncUpNodesRequest syncU
 
         cleanUpLocalCache(runningDeployModelTasks);
         cleanUpLocalCacheFiles();
+        if (mlFeatureEnabledSetting.isKeyRefreshEnabled()) {
+            encryptor.refreshMasterKey(ActionListener.wrap(r -> {
+                if (r) {
+                    log.debug("Refresh ML key completed.");
+                    return;
+                }
+            }, e -> { log.error("Failed to refresh ML key", e); }));
+        }
 
         return new MLSyncUpNodeResponse(
             clusterService.localNode(),

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

@@ -1156,7 +1156,8 @@ public List<Setting<?>> getSettings() {
                 MLCommonsSettings.ML_COMMONS_MCP_CONNECTOR_ENABLED,
                 MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED,
                 MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED,
-                MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED
+                MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED,
+                MLCommonsSettings.ML_COMMONS_KEY_REFRESH_ENABLED
             );
         return settings;
     }

plugin/src/test/java/org/opensearch/ml/action/syncup/TransportSyncUpOnNodeActionTests.java

@@ -11,6 +11,7 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -57,13 +58,15 @@
 import org.opensearch.ml.common.MLTask;
 import org.opensearch.ml.common.MLTaskType;
 import org.opensearch.ml.common.model.MLModelState;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.transport.sync.MLSyncUpInput;
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodeRequest;
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodeResponse;
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodesRequest;
 import org.opensearch.ml.common.transport.sync.MLSyncUpNodesResponse;
 import org.opensearch.ml.engine.MLEngine;
 import org.opensearch.ml.engine.ModelHelper;
+import org.opensearch.ml.engine.encryptor.EncryptorImpl;
 import org.opensearch.ml.model.MLModelCacheHelper;
 import org.opensearch.ml.model.MLModelManager;
 import org.opensearch.ml.task.MLTaskCache;
@@ -119,11 +122,19 @@ public class TransportSyncUpOnNodeActionTests extends OpenSearchTestCase {
     @Mock
     private MLModelCacheHelper mlModelCacheHelper;
 
+    @Mock
+    private EncryptorImpl encryptor;
+
+    @Mock
+    MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
     @Before
     public void setup() throws IOException {
         MockitoAnnotations.openMocks(this);
         mockSettings(true);
         when(clusterService.getClusterName()).thenReturn(new ClusterName("Local Cluster"));
+        when(mlFeatureEnabledSetting.isKeyRefreshEnabled()).thenReturn(false);
+        doNothing().when(encryptor).refreshMasterKey(any());
         action = new TransportSyncUpOnNodeAction(
             transportService,
             settings,
@@ -136,7 +147,9 @@ public void setup() throws IOException {
             client,
             xContentRegistry,
             mlEngine,
-            mlModelCacheHelper
+            mlModelCacheHelper,
+            encryptor,
+            mlFeatureEnabledSetting
         );
         runningDeployModelTasks = new HashMap<>();
         runningDeployModelTasks.put("model1", ImmutableSet.of("node1"));

plugin/src/test/java/org/opensearch/ml/settings/MLFeatureEnabledSettingTests.java

@@ -14,6 +14,7 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENT_FRAMEWORK_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_KEY_REFRESH_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_LOCAL_MODEL_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED;
@@ -70,7 +71,8 @@ public void setUp() {
                             ML_COMMONS_MCP_SERVER_ENABLED,
                             ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED,
                             ML_COMMONS_METRIC_COLLECTION_ENABLED,
-                            ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED
+                            ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED,
+                            ML_COMMONS_KEY_REFRESH_ENABLED
                         )
                 )
             );