From 2c05b58c74c7c05897a57ba813af09712f742931 Mon Sep 17 00:00:00 2001 From: Yaliang Wu Date: Wed, 12 Jul 2023 03:29:00 -0700 Subject: [PATCH] fix init master key bug Signed-off-by: Yaliang Wu --- .../ml/engine/encryptor/EncryptorImpl.java | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java index 83502fb85c..f724b4698a 100644 --- a/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java +++ b/ml-algorithms/src/main/java/org/opensearch/ml/engine/encryptor/EncryptorImpl.java @@ -27,6 +27,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicReference; +import static java.util.concurrent.TimeUnit.SECONDS; import static org.opensearch.ml.common.CommonValue.MASTER_KEY; import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX; @@ -111,7 +112,7 @@ private void initMasterKey() { client.get(getRequest, new LatchedActionListener(ActionListener.wrap(r -> { if (r.isExists()) { String masterKey = (String) r.getSourceAsMap().get(MASTER_KEY); - setMasterKey(masterKey); + this.masterKey = masterKey; } else { exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet")); } @@ -122,6 +123,13 @@ private void initMasterKey() { } } else { exceptionRef.set(new ResourceNotFoundException("ML encryption master key not initialized yet")); + latch.countDown(); + } + + try { + latch.await(5, SECONDS); + } catch (InterruptedException e) { + throw new IllegalStateException(e); } if (exceptionRef.get() != null) { @@ -132,5 +140,8 @@ private void initMasterKey() { throw new MLException(exceptionRef.get()); } } + if (masterKey == null) { + throw new ResourceNotFoundException("ML encryption master key not initialized yet"); + } } }