Fix connector can use subfield of secrectArn to bypass extraction of …

…credential from AWS (opensearch-project#2320) (opensearch-project#2321) * Fix connector can use subfield of secrectArn to bypass extraction of credential from AWS * Forbid all kinds of ARN subfield * Adjust annotation --------- Signed-off-by: Sicheng Song <[email protected]>
dhrubo-os · May 17, 2024 · 2ff5e96 · 2ff5e96
1 parent ca53e59
commit 2ff5e96
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 9 deletions.
diff --git a/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java b/...on/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInput.java
@@ -25,6 +25,8 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;
@@ -148,7 +150,17 @@ public static MLCreateConnectorInput parse(XContentParser parser, boolean update
                     parameters = getParameterMap(parser.map());
                     break;
                 case CONNECTOR_CREDENTIAL_FIELD:
-                    credential = parser.mapStrings();
+                    // We need to filter out any key string that is trying to imitate the subfield of any kind of ARN of the credential map
+                    credential = new HashMap<>();
+                    Map<String, String> credentialKeyToAdd = parser.mapStrings();
+                    Pattern pattern = Pattern.compile("[a-zA-Z]+Arn\\.");
+                    for (String key : credentialKeyToAdd.keySet()) {
+                        Matcher matcher = pattern.matcher(key);
+                        boolean matchFound = matcher.find();
+                        if (!matchFound) {
+                            credential.put(key, credentialKeyToAdd.get(key));
+                        }
+                    }
                     break;
                 case CONNECTOR_ACTIONS_FIELD:
                     actions = new ArrayList<>();

diff --git a/...c/test/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInputTests.java b/...c/test/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorInputTests.java
@@ -5,13 +5,6 @@
 
 package org.opensearch.ml.common.transport.connector;
 
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.function.Consumer;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -20,9 +13,9 @@
 import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.xcontent.LoggingDeprecationHandler;
-import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.common.io.stream.StreamInput;
+import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -34,6 +27,15 @@
 import org.opensearch.ml.common.connector.MLPreProcessFunction;
 import org.opensearch.search.SearchModule;
 
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
@@ -190,6 +192,28 @@ public void testParse_ArrayParameter() throws Exception {
         });
     }
 
+    @Test
+    public void testParse_SecretArnPrefix() throws Exception {
+        String expectedInputStr = "{\"name\":\"test_connector_name\"," +
+                "\"description\":\"this is a test connector\",\"version\":\"1\",\"protocol\":\"http\"," +
+                "\"parameters\":{\"input\":[\"test input value\"]},\"credential\":{\"key\":\"test_key_value\"," +
+                "\"secretArn\":\"test_secretArn_value\", \"secretArn.key\":\"test_key_value\"," +
+                "\"roleArn\":\"test_roleArn_value\", \"roleArn.subfield\":\"test_subfield_value\",\"test_Arn_test\":\"test_value\"}," +
+                "\"actions\":[{\"action_type\":\"PREDICT\",\"method\":\"POST\",\"url\":\"https://test.com\"," +
+                "\"headers\":{\"api_key\":\"${credential.key}\"}," +
+                "\"request_body\":\"{\\\"input\\\": \\\"${parameters.input}\\\"}\"," +
+                "\"pre_process_function\":\"connector.pre_process.openai.embedding\"," +
+                "\"post_process_function\":\"connector.post_process.openai.embedding\"}]," +
+                "\"backend_roles\":[\"role1\",\"role2\"],\"add_all_backend_roles\":false," +
+                "\"access_mode\":\"PUBLIC\"}";
+        HashSet<String> expectedCredentialKeys = new HashSet<>(Arrays.asList("key", "secretArn", "roleArn","test_Arn_test"));
+        testParseFromJsonString(expectedInputStr, parsedInput -> {
+            assertEquals(expectedCredentialKeys, parsedInput.getCredential().keySet());
+            assertEquals("test_secretArn_value", parsedInput.getCredential().get("secretArn"));
+            assertEquals("test_roleArn_value", parsedInput.getCredential().get("roleArn"));
+        });
+    }
+
     @Test
     public void testParseWithDryRun() throws Exception {
         String expectedInputStrWithDryRun = "{\"dry_run\":true}";