Skip to content

Commit 20eb9bc

Browse files
b4sjoob4sjoo
authored
Fix 403 during execute connector (#4414)
1 parent addda96 commit 20eb9bc

File tree

3 files changed

+73
-16
lines changed

3 files changed

+73
-16
lines changed

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java‎

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1044,51 +1044,60 @@ public static Map<String, Object> createMemoryParams(
10441044
memoryParams.put(TENANT_ID_FIELD, mlAgent.getTenantId());
10451045
}
10461046
if (requestParameters != null) {
1047-
// Check if parameters are wrapped in remote_agent_memory_configuration
1048-
String remoteMemoryConfigStr = requestParameters.get("remote_agent_memory_configuration");
1049-
if (!Strings.isNullOrEmpty(remoteMemoryConfigStr)) {
1050-
// Parse the remote_agent_memory_configuration JSON
1047+
// Check if parameters are wrapped in memory_configuration
1048+
String memoryConfigStr = requestParameters.get("memory_configuration");
1049+
if (!Strings.isNullOrEmpty(memoryConfigStr)) {
1050+
// Parse the memory_configuration JSON
10511051
try (
10521052
XContentParser parser = JsonXContent.jsonXContent
1053-
.createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, remoteMemoryConfigStr)
1053+
.createParser(NamedXContentRegistry.EMPTY, LoggingDeprecationHandler.INSTANCE, memoryConfigStr)
10541054
) {
10551055
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
1056-
Map<String, Object> remoteMemoryConfig = parser.map();
1056+
Map<String, Object> memoryConfig = parser.map();
10571057

10581058
// Extract memory_container_id
1059-
String memoryContainerIdParam = (String) remoteMemoryConfig.get(MEMORY_CONTAINER_ID_FIELD);
1059+
String memoryContainerIdParam = (String) memoryConfig.get(MEMORY_CONTAINER_ID_FIELD);
10601060
if (!Strings.isNullOrEmpty(memoryContainerIdParam)) {
10611061
memoryParams.put(MEMORY_CONTAINER_ID_FIELD, memoryContainerIdParam);
10621062
}
10631063

10641064
// Extract endpoint
1065-
String endpointParam = (String) remoteMemoryConfig.get(ENDPOINT_FIELD);
1065+
String endpointParam = (String) memoryConfig.get(ENDPOINT_FIELD);
10661066
if (!Strings.isNullOrEmpty(endpointParam)) {
10671067
memoryParams.put(ENDPOINT_FIELD, endpointParam);
10681068
}
10691069

10701070
// Extract region
1071-
String regionParam = (String) remoteMemoryConfig.get(HttpConnector.REGION_FIELD);
1071+
String regionParam = (String) memoryConfig.get(HttpConnector.REGION_FIELD);
10721072
if (!Strings.isNullOrEmpty(regionParam)) {
10731073
memoryParams.put(HttpConnector.REGION_FIELD, regionParam);
10741074
}
10751075

10761076
// Extract credential
1077-
Object credentialObj = remoteMemoryConfig.get(CREDENTIAL_FIELD);
1077+
Object credentialObj = memoryConfig.get(CREDENTIAL_FIELD);
10781078
if (credentialObj instanceof Map) {
10791079
Map<String, String> credential = (Map<String, String>) credentialObj;
10801080
if (!credential.isEmpty()) {
10811081
memoryParams.put(CREDENTIAL_FIELD, credential);
10821082
}
10831083
}
10841084

1085+
// Check for direct roleArn field - if present, override credential map
1086+
String roleArnParam = (String) memoryConfig.get("roleArn");
1087+
if (!Strings.isNullOrEmpty(roleArnParam)) {
1088+
// Override credential with roleArn
1089+
Map<String, String> roleArnCredential = new HashMap<>();
1090+
roleArnCredential.put("roleArn", roleArnParam);
1091+
memoryParams.put(CREDENTIAL_FIELD, roleArnCredential);
1092+
}
1093+
10851094
// Extract user_id if provided
1086-
String userIdParam = (String) remoteMemoryConfig.get("user_id");
1095+
String userIdParam = (String) memoryConfig.get("user_id");
10871096
if (!Strings.isNullOrEmpty(userIdParam)) {
10881097
memoryParams.put("user_id", userIdParam);
10891098
}
10901099
} catch (Exception e) {
1091-
log.error("Failed to parse remote_agent_memory_configuration", e);
1100+
log.error("Failed to parse memory_configuration", e);
10921101
}
10931102
}
10941103
memoryParams.put(TENANT_ID_FIELD, mlAgent.getTenantId());

‎ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemory.java‎

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,6 +1032,9 @@ private Connector createInlineConnector(String endpoint, String region, Map<Stri
10321032
credentials.putAll(credential);
10331033
}
10341034

1035+
// Extract tenant ID from role ARN if applicable
1036+
String tenantId = extractTenantIdFromRoleArn(serviceName, credentials);
1037+
10351038
// Create Memory Container API actions
10361039
List<ConnectorAction> actions = createMemoryContainerActions();
10371040

@@ -1046,6 +1049,7 @@ private Connector createInlineConnector(String endpoint, String region, Map<Stri
10461049
.parameters(parameters)
10471050
.credential(credentials)
10481051
.actions(actions)
1052+
.tenantId(tenantId)
10491053
.build();
10501054
} else {
10511055
// Use HttpConnector for plain HTTP
@@ -1056,6 +1060,7 @@ private Connector createInlineConnector(String endpoint, String region, Map<Stri
10561060
.parameters(parameters)
10571061
.credential(credentials.isEmpty() ? null : credentials)
10581062
.actions(actions)
1063+
.tenantId(null)
10591064
.build();
10601065
}
10611066

@@ -1087,8 +1092,8 @@ private Connector createInlineConnector(String endpoint, String region, Map<Stri
10871092
connector
10881093
.decrypt(
10891094
ConnectorAction.ActionType.EXECUTE.name(),
1090-
(cred, tenantId) -> cred, // No-op function - credentials are already plaintext
1091-
null // No tenant ID for inline connectors
1095+
(cred, tenant) -> cred, // No-op function - credentials are already plaintext
1096+
tenantId
10921097
);
10931098

10941099
return connector;
@@ -1211,6 +1216,48 @@ private String extractServiceName(String endpoint) {
12111216
return "aoss";
12121217
}
12131218

1219+
/**
1220+
* Extract dummy tenant ID from role ARN for AOSS services.
1221+
* Essentially we only need the front part (account ID) as client ID when in AOSS
1222+
*
1223+
* @param serviceName The AWS service name (aoss or es)
1224+
* @param credential The credential map that may contain roleArn
1225+
* @return Tenant ID in format "account:role" for AOSS, null for ES or if roleArn not present
1226+
*
1227+
* Example:
1228+
* - Input: serviceName="aoss", roleArn="arn:aws:iam::123456789012:role/role-name"
1229+
* - Output: "123456789012:role"
1230+
*/
1231+
private String extractTenantIdFromRoleArn(String serviceName, Map<String, String> credential) {
1232+
// Return null for ES service
1233+
if (!"aoss".equals(serviceName)) {
1234+
return null;
1235+
}
1236+
1237+
// Check if credential map exists and contains roleArn
1238+
if (credential == null || !credential.containsKey("roleArn")) {
1239+
return null;
1240+
}
1241+
1242+
String roleArn = credential.get("roleArn");
1243+
if (Strings.isNullOrEmpty(roleArn)) {
1244+
return null;
1245+
}
1246+
1247+
// Expected format: arn:aws:iam::{account}:role/{role-name}
1248+
try {
1249+
String[] parts = roleArn.split(":");
1250+
if (parts.length >= 5 && "role".equals(parts[4])) {
1251+
String account = parts[3];
1252+
return account + ":role";
1253+
}
1254+
} catch (Exception e) {
1255+
log.error("Failed to parse roleArn: {}", roleArn, e);
1256+
}
1257+
1258+
return null;
1259+
}
1260+
12141261
/**
12151262
* Validate endpoint URL
12161263
*/

‎plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,11 @@ private void executeWithConnector(
107107
ActionListener<MLTaskResponse> listener,
108108
boolean decryptWithEncryptor
109109
) {
110+
String connectorTenantId = connector.getTenantId();
110111
if (decryptWithEncryptor) {
111-
connector.decrypt(action, (credential, tenantId) -> encryptor.decrypt(credential, null), null);
112+
connector.decrypt(action, (credential, tenantId) -> encryptor.decrypt(credential, tenantId), connectorTenantId);
112113
} else {
113-
connector.decrypt(action, (credential, tenantId) -> credential, null);
114+
connector.decrypt(action, (credential, tenantId) -> credential, connectorTenantId);
114115
}
115116
RemoteConnectorExecutor connectorExecutor = MLEngineClassLoader.initInstance(connector.getProtocol(), connector, Connector.class);
116117
connectorExecutor.setConnectorPrivateIpEnabled(mlFeatureEnabledSetting.isConnectorPrivateIpEnabled());

0 commit comments

Comments
 (0)