Adds a session token to AWS Credentials (opensearch-project#6103)

* Adds session token for aws connection Signed-off-by: Bandini Bhopi <[email protected]> * Adds changelog Signed-off-by: Bandini Bhopi <[email protected]> --------- Signed-off-by: Bandini Bhopi <[email protected]>
xinruiba · Mar 11, 2024 · 6d882c9 · 6d882c9
1 parent ead2947
commit 6d882c9
Show file tree

Hide file tree

Showing 8 changed files with 156 additions and 37 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -27,6 +27,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - [Multiple Datasource] Expose a few properties for customize the appearance of the data source selector component ([#6057](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6057))
 - [Multiple Datasource] Create data source menu component able to be mount to nav bar ([#6082](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6082))
 - [Multiple Datasource] Handle form values(request payload) if the selected type is available in the authentication registry ([#6049](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6049))
+- [Multiple Datasource] Adds a session token to AWS credentials ([#6103](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6103))
 - [Multiple Datasource] Add Vega support to MDS by specifying a data source name in the Vega spec ([#5975](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/5975))
 - [Workspace] Consume workspace id in saved object client ([#6014](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6014))
 - [Multiple Datasource] Export DataSourcePluginRequestContext at top level for plugins to use ([#6108](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6108))

diff --git a/src/plugins/data_source/server/client/configure_client.test.mocks.ts b/src/plugins/data_source/server/client/configure_client.test.mocks.ts
@@ -21,3 +21,12 @@ export const authRegistryCredentialProviderMock = jest.fn();
 jest.doMock('../util/credential_provider', () => ({
   authRegistryCredentialProvider: authRegistryCredentialProviderMock,
 }));
+
+export const CredentialsMock = jest.fn();
+jest.doMock('aws-sdk', () => {
+  const actual = jest.requireActual('aws-sdk');
+  return {
+    ...actual,
+    Credentials: CredentialsMock,
+  };
+});
diff --git a/src/plugins/data_source/server/client/configure_client.test.ts b/src/plugins/data_source/server/client/configure_client.test.ts
@@ -17,6 +17,7 @@ import {
   ClientMock,
   parseClientOptionsMock,
   authRegistryCredentialProviderMock,
+  CredentialsMock,
 } from './configure_client.test.mocks';
 import { OpenSearchClientPoolSetup } from './client_pool';
 import { configureClient } from './configure_client';
@@ -48,6 +49,17 @@ describe('configureClient', () => {
   let customApiSchemaRegistry: CustomApiSchemaRegistry;
   let authenticationMethodRegistery: jest.Mocked<IAuthenticationMethodRegistery>;
 
+  const customAuthContent = {
+    region: 'us-east-1',
+    roleARN: 'test-role',
+  };
+
+  const authMethod: AuthenticationMethod = {
+    name: 'typeA',
+    authType: AuthType.SigV4,
+    credentialProvider: jest.fn(),
+  };
+
   beforeEach(() => {
     dsClient = opensearchClientMock.createInternalClient();
     logger = loggingSystemMock.createLogger();
@@ -110,10 +122,12 @@ describe('configureClient', () => {
     };
 
     ClientMock.mockImplementation(() => dsClient);
+    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
   });
 
   afterEach(() => {
     ClientMock.mockReset();
+    CredentialsMock.mockReset();
   });
 
   test('configure client with auth.type == no_auth, will call new Client() to create client', async () => {
@@ -251,12 +265,7 @@ describe('configureClient', () => {
     expect(decodeAndDecryptSpy).toHaveBeenCalledTimes(1);
   });
 
-  test('configureClient should retunrn client from authentication registery if method present in registry', async () => {
-    const name = 'typeA';
-    const customAuthContent = {
-      region: 'us-east-1',
-      roleARN: 'test-role',
-    };
+  test('configureClient should return client if authentication method from registry provides credentials', async () => {
     savedObjectsMock.get.mockReset().mockResolvedValueOnce({
       id: DATA_SOURCE_ID,
       type: DATA_SOURCE_SAVED_OBJECT_TYPE,
@@ -269,12 +278,6 @@ describe('configureClient', () => {
       },
       references: [],
     });
-    const authMethod: AuthenticationMethod = {
-      name,
-      authType: AuthType.SigV4,
-      credentialProvider: jest.fn(),
-    };
-    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
 
     authRegistryCredentialProviderMock.mockReturnValue({
       credential: sigV4AuthContent,
@@ -291,5 +294,46 @@ describe('configureClient', () => {
     expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
     expect(ClientMock).toHaveBeenCalledTimes(1);
     expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toBeCalledWith({
+      accessKeyId: sigV4AuthContent.accessKey,
+      secretAccessKey: sigV4AuthContent.secretKey,
+    });
+  });
+
+  test('When credential provider from auth registry returns session token, credentials should contains session token', async () => {
+    const mockCredentials = { ...sigV4AuthContent, sessionToken: 'sessionToken' };
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: AuthType.SigV4,
+          credentials: customAuthContent,
+        },
+      },
+      references: [],
+    });
+
+    authRegistryCredentialProviderMock.mockReturnValue({
+      credential: mockCredentials,
+      type: AuthType.SigV4,
+    });
+
+    await configureClient(
+      { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
+      clientPoolSetup,
+      config,
+      logger
+    );
+
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toBeCalledWith({
+      accessKeyId: mockCredentials.accessKey,
+      secretAccessKey: mockCredentials.secretKey,
+      sessionToken: mockCredentials.sessionToken,
+    });
   });
 });
diff --git a/src/plugins/data_source/server/client/configure_client.ts b/src/plugins/data_source/server/client/configure_client.ts
@@ -26,6 +26,7 @@ import {
   getCredential,
   getDataSource,
   generateCacheKey,
+  getSigV4Credentials,
 } from './configure_client_utils';
 import { IAuthenticationMethodRegistery } from '../auth_registry';
 import { authRegistryCredentialProvider } from '../util/credential_provider';
@@ -199,11 +200,12 @@ const getBasicAuthClient = (
 };
 
 const getAWSClient = (credential: SigV4Content, clientOptions: ClientOptions): Client => {
-  const { accessKey, secretKey, region, service } = credential;
+  const { accessKey, secretKey, region, service, sessionToken } = credential;
+  const sigv4Credentials = getSigV4Credentials(accessKey, secretKey, sessionToken);
 
   const credentialProvider = (): Promise<Credentials> => {
     return new Promise((resolve) => {
-      resolve(new Credentials({ accessKeyId: accessKey, secretAccessKey: secretKey }));
+      resolve(sigv4Credentials);
     });
   };
 

diff --git a/src/plugins/data_source/server/client/configure_client_utils.ts b/src/plugins/data_source/server/client/configure_client_utils.ts
@@ -5,6 +5,7 @@
 
 import { Client } from '@opensearch-project/opensearch';
 import { Client as LegacyClient } from 'elasticsearch';
+import { Credentials } from 'aws-sdk';
 import { SavedObjectsClientContract } from '../../../../../src/core/server';
 import { DATA_SOURCE_SAVED_OBJECT_TYPE } from '../../common';
 import {
@@ -145,3 +146,21 @@ export const generateCacheKey = (dataSourceAttr: DataSourceAttributes, dataSourc
 
   return key;
 };
+
+export const getSigV4Credentials = (
+  accessKeyId: string,
+  secretAccessKey: string,
+  sessionToken?: string
+): Credentials => {
+  let sigv4Credentials: Credentials;
+  if (sessionToken) {
+    sigv4Credentials = new Credentials({
+      accessKeyId,
+      secretAccessKey,
+      sessionToken,
+    });
+  } else {
+    sigv4Credentials = new Credentials({ accessKeyId, secretAccessKey });
+  }
+  return sigv4Credentials;
+};
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.test.mocks.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.test.mocks.ts
@@ -16,8 +16,3 @@ export const parseClientOptionsMock = jest.fn();
 jest.doMock('./client_config', () => ({
   parseClientOptions: parseClientOptionsMock,
 }));
-
-export const authRegistryCredentialProviderMock = jest.fn();
-jest.doMock('../util/credential_provider', () => ({
-  authRegistryCredentialProvider: authRegistryCredentialProviderMock,
-}));
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.test.ts
@@ -13,11 +13,11 @@ import { CryptographyServiceSetup } from '../cryptography_service';
 import { DataSourceClientParams, LegacyClientCallAPIParams, AuthenticationMethod } from '../types';
 import { OpenSearchClientPoolSetup } from '../client';
 import { ConfigOptions } from 'elasticsearch';
+import { ClientMock, parseClientOptionsMock } from './configure_legacy_client.test.mocks';
 import {
-  ClientMock,
-  parseClientOptionsMock,
   authRegistryCredentialProviderMock,
-} from './configure_legacy_client.test.mocks';
+  CredentialsMock,
+} from '../client/./configure_client.test.mocks';
 import { configureLegacyClient } from './configure_legacy_client';
 import { CustomApiSchemaRegistry } from '../schema_registry';
 import { IAuthenticationMethodRegistery } from '../auth_registry';
@@ -47,6 +47,17 @@ describe('configureLegacyClient', () => {
 
   const mockResponse = { data: 'ping' };
 
+  const customAuthContent = {
+    region: 'us-east-1',
+    roleARN: 'test-role',
+  };
+
+  const authMethod: AuthenticationMethod = {
+    name: 'typeA',
+    authType: AuthType.SigV4,
+    credentialProvider: jest.fn(),
+  };
+
   beforeEach(() => {
     mockOpenSearchClientInstance = {
       close: jest.fn(),
@@ -119,10 +130,13 @@ describe('configureLegacyClient', () => {
         response: mockResponse,
       });
     });
+
+    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
   });
 
   afterEach(() => {
     ClientMock.mockReset();
+    CredentialsMock.mockReset();
     jest.resetAllMocks();
   });
 
@@ -263,12 +277,7 @@ describe('configureLegacyClient', () => {
     expect(mockOpenSearchClientInstance.ping).toHaveBeenLastCalledWith(mockParams);
   });
 
-  test('configureLegacyClient should retunrn client from authentication registery if method present in registry', async () => {
-    const name = 'typeA';
-    const customAuthContent = {
-      region: 'us-east-1',
-      roleARN: 'test-role',
-    };
+  test('configureLegacyClient should return client if authentication method from registry provides credentials', async () => {
     savedObjectsMock.get.mockReset().mockResolvedValueOnce({
       id: DATA_SOURCE_ID,
       type: DATA_SOURCE_SAVED_OBJECT_TYPE,
@@ -281,12 +290,6 @@ describe('configureLegacyClient', () => {
       },
       references: [],
     });
-    const authMethod: AuthenticationMethod = {
-      name,
-      authType: AuthType.SigV4,
-      credentialProvider: jest.fn(),
-    };
-    authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
 
     authRegistryCredentialProviderMock.mockReturnValue({
       credential: sigV4AuthContent,
@@ -304,5 +307,49 @@ describe('configureLegacyClient', () => {
     expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
     expect(ClientMock).toHaveBeenCalledTimes(1);
     expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toBeCalledWith({
+      accessKeyId: sigV4AuthContent.accessKey,
+      secretAccessKey: sigV4AuthContent.secretKey,
+    });
+  });
+
+  test('When credential provider from auth registry returns session token, credentials should contains session token', async () => {
+    const mockCredentials = { ...sigV4AuthContent, sessionToken: 'sessionToken' };
+    savedObjectsMock.get.mockReset().mockResolvedValueOnce({
+      id: DATA_SOURCE_ID,
+      type: DATA_SOURCE_SAVED_OBJECT_TYPE,
+      attributes: {
+        ...dataSourceAttr,
+        auth: {
+          type: AuthType.SigV4,
+          credentials: customAuthContent,
+        },
+      },
+      references: [],
+    });
+
+    authRegistryCredentialProviderMock.mockReturnValue({
+      credential: mockCredentials,
+      type: AuthType.SigV4,
+    });
+
+    await configureLegacyClient(
+      { ...dataSourceClientParams, authRegistry: authenticationMethodRegistery },
+      callApiParams,
+      clientPoolSetup,
+      config,
+      logger
+    );
+    expect(authRegistryCredentialProviderMock).toHaveBeenCalled();
+    expect(authenticationMethodRegistery.getAuthenticationMethod).toHaveBeenCalledTimes(1);
+    expect(ClientMock).toHaveBeenCalledTimes(1);
+    expect(savedObjectsMock.get).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toHaveBeenCalledTimes(1);
+    expect(CredentialsMock).toBeCalledWith({
+      accessKeyId: mockCredentials.accessKey,
+      secretAccessKey: mockCredentials.secretKey,
+      sessionToken: mockCredentials.sessionToken,
+    });
   });
 });
diff --git a/src/plugins/data_source/server/legacy/configure_legacy_client.ts b/src/plugins/data_source/server/legacy/configure_legacy_client.ts
@@ -5,7 +5,7 @@
 
 import { Client } from '@opensearch-project/opensearch';
 import { Client as LegacyClient, ConfigOptions } from 'elasticsearch';
-import { Credentials, Config } from 'aws-sdk';
+import { Config } from 'aws-sdk';
 import { get } from 'lodash';
 import HttpAmazonESConnector from 'http-aws-es';
 import {
@@ -34,6 +34,7 @@ import {
   getCredential,
   getDataSource,
   generateCacheKey,
+  getSigV4Credentials,
 } from '../client/configure_client_utils';
 import { IAuthenticationMethodRegistery } from '../auth_registry';
 import { authRegistryCredentialProvider } from '../util/credential_provider';
@@ -230,12 +231,13 @@ const getBasicAuthClient = async (
 };
 
 const getAWSClient = (credential: SigV4Content, clientOptions: ConfigOptions): LegacyClient => {
-  const { accessKey, secretKey, region, service } = credential;
+  const { accessKey, secretKey, region, service, sessionToken } = credential;
+  const credentials = getSigV4Credentials(accessKey, secretKey, sessionToken);
   const client = new LegacyClient({
     connectionClass: HttpAmazonESConnector,
     awsConfig: new Config({
       region,
-      credentials: new Credentials({ accessKeyId: accessKey, secretAccessKey: secretKey }),
+      credentials,
     }),
     service,
     ...clientOptions,