From 11a35717998ecadb24c911ea738556d63b33b3ba Mon Sep 17 00:00:00 2001 From: "opensearch-trigger-bot[bot]" <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com> Date: Wed, 12 Jul 2023 03:33:41 -0700 Subject: [PATCH] fix init master key bug (#1094) (#1096) Signed-off-by: Yaliang Wu (cherry picked from commit b782f29e79533eea6e82d53a1eba52e7deb40682) Co-authored-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"); + } } }