fix: Make lazy-loading of secrets atomic and handle updateValue with unloaded secrets

Use ConcurrentMap.compute() in loadSecretIfNeeded() to ensure the
sentinel check and secret retrieval are performed atomically, avoiding
race conditions with concurrent access.

Add loadSecretIfNeeded() call at the beginning of updateValue() to
ensure secrets are loaded before update logic runs, preventing the
NOT_LOADED_SENTINEL from being treated as a plain value store.

Signed-off-by: Sumit Bhattacharya <sumit4739@gmail.com>

Author: BhattacharyaSumit

data-prepper-plugins/aws-plugin/src/main/java/org/opensearch/dataprepper/plugins/aws/AwsSecretsSupplier.java

@@ -123,18 +123,21 @@ public Object retrieveValue(String secretId) {
 
     /**
      * Loads a secret if it was skipped on start (lazy-loading).
-     * 
+     * Uses {@link ConcurrentMap#compute} to ensure atomicity of the sentinel check and refresh.
+     *
      * @param secretId The secret configuration ID
      * @return The loaded secret value
      */
     private Object loadSecretIfNeeded(String secretId) {
-        Object value = secretIdToValue.get(secretId);
-        if (value == NOT_LOADED_SENTINEL) {
-            LOG.info("Secret {} was not loaded on start, loading now on first access.", secretId);
-            refresh(secretId);
-            value = secretIdToValue.get(secretId);
-        }
-        return value;
+        return secretIdToValue.compute(secretId, (key, currentValue) -> {
+            if (currentValue == NOT_LOADED_SENTINEL) {
+                LOG.info("Secret {} was not loaded on start, loading now on first access.", key);
+                final AwsSecretManagerConfiguration config = awsSecretManagerConfigurationMap.get(key);
+                final SecretsManagerClient client = secretsManagerClientMap.get(key);
+                return retrieveSecretsFromSecretManager(config, client);
+            }
+            return currentValue;
+        });
     }
 
 
@@ -184,6 +187,8 @@ public String updateValue(String secretId, Object newValue) {
 
     @Override
     public String updateValue(String secretId, String keyToUpdate, Object newValue) {
+        // Ensure the secret is loaded before attempting to update
+        loadSecretIfNeeded(secretId);
         Object currentSecretStore = secretIdToValue.get(secretId);
         if (currentSecretStore instanceof Map) {
             if (keyToUpdate == null) {

data-prepper-plugins/aws-plugin/src/test/java/org/opensearch/dataprepper/plugins/aws/AwsSecretsSupplierLazyLoadTest.java

@@ -20,12 +20,15 @@
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest;
 import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse;
+import software.amazon.awssdk.services.secretsmanager.model.PutSecretValueRequest;
+import software.amazon.awssdk.services.secretsmanager.model.PutSecretValueResponse;
 
 import java.util.Map;
 import java.util.UUID;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -61,6 +64,12 @@ class AwsSecretsSupplierLazyLoadTest {
     @Mock
     private GetSecretValueResponse getSecretValueResponse;
 
+    @Mock
+    private PutSecretValueRequest putSecretValueRequest;
+
+    @Mock
+    private PutSecretValueResponse putSecretValueResponse;
+
     @Mock
     private AwsCredentialsSupplier awsCredentialsSupplier;
 
@@ -131,4 +140,37 @@ void testSecretWithSkipValidationOnStartFalse_LoadsAtConstruction() throws JsonP
         verify(secretsManagerClient, times(1)).getSecretValue(eq(getSecretValueRequest));
         assertThat(value, equalTo(testValue));
     }
+
+    @Test
+    void testUpdateValue_withSkipValidationOnStart_loadsSecretBeforeUpdate() throws JsonProcessingException {
+        // Given: Secret configured with skip_validation_on_start=true
+        when(awsSecretPluginConfig.getAwsSecretManagerConfigurationMap()).thenReturn(
+                Map.of(testSecretId, awsSecretManagerConfiguration)
+        );
+        when(awsSecretManagerConfiguration.isSkipValidationOnStart()).thenReturn(true);
+        when(awsSecretManagerConfiguration.createSecretManagerClient(awsCredentialsSupplier)).thenReturn(secretsManagerClient);
+        when(awsSecretManagerConfiguration.createGetSecretValueRequest()).thenReturn(getSecretValueRequest);
+        when(secretValueDecoder.decode(eq(getSecretValueResponse))).thenReturn(objectMapper.writeValueAsString(
+                Map.of(testKey, testValue)
+        ));
+        when(secretsManagerClient.getSecretValue(eq(getSecretValueRequest))).thenReturn(getSecretValueResponse);
+        when(awsSecretManagerConfiguration.putSecretValueRequest(any())).thenReturn(putSecretValueRequest);
+        when(secretsManagerClient.putSecretValue(eq(putSecretValueRequest))).thenReturn(putSecretValueResponse);
+        final String newVersionId = UUID.randomUUID().toString();
+        when(putSecretValueResponse.versionId()).thenReturn(newVersionId);
+
+        final AwsSecretsSupplier supplier = new AwsSecretsSupplier(
+                secretValueDecoder, awsSecretPluginConfig, objectMapper, awsCredentialsSupplier
+        );
+
+        // Then: Secret is NOT retrieved at construction time
+        verify(secretsManagerClient, never()).getSecretValue(eq(getSecretValueRequest));
+
+        // When: updateValue is called before any retrieveValue
+        final String versionId = supplier.updateValue(testSecretId, testKey, "newValue");
+
+        // Then: Secret was loaded on-demand and update succeeded
+        verify(secretsManagerClient, times(1)).getSecretValue(eq(getSecretValueRequest));
+        assertThat(versionId, equalTo(newVersionId));
+    }
 }