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 index 98d4811086fd..d1135531da9c 100644 --- 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 index 27876f6ea527..9850d8773371 100644 --- 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; + 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 index d12688df0707..20d3d3773960 100644 --- 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 => { 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 index 4405f8f14559..9392b21a24b8 100644 --- 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-next'; 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 index 2f91e757fd28..e6c1b3363896 100644 --- 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 index ebe356c58561..5685392dbc10 100644 --- 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 index 42dbd4aa2772..05bd5dbc1d27 100644 --- 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-next'; 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,