Skip to content

Commit

Permalink
Adds a session token to AWS Credentials (opensearch-project#6103) (op…
Browse files Browse the repository at this point in the history 
…ensearch-project#6112)

* Adds session token for aws connection

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>

* Adds changelog

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>

---------

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>
(cherry picked from commit 6d882c9)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
  • Loading branch information
@opensearch-trigger-bot @github-actions
opensearch-trigger-bot[bot] and github-actions[bot] authored Mar 11, 2024
1 parent 52e639f commit 051087d
Show file tree
Hide file tree
Showing 7 changed files with 155 additions and 37 deletions.
9 changes: 9 additions & 0 deletions src/plugins/data_source/server/client/configure_client.test.mocks.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
});
68 changes: 56 additions & 12 deletions src/plugins/data_source/server/client/configure_client.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ import {
ClientMock,
parseClientOptionsMock,
authRegistryCredentialProviderMock,
CredentialsMock,
} from './configure_client.test.mocks';
import { OpenSearchClientPoolSetup } from './client_pool';
import { configureClient } from './configure_client';
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
});
});
});
6 changes: 4 additions & 2 deletions src/plugins/data_source/server/client/configure_client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import {
getCredential,
getDataSource,
generateCacheKey,
getSigV4Credentials,
} from './configure_client_utils';
import { IAuthenticationMethodRegistery } from '../auth_registry';
import { authRegistryCredentialProvider } from '../util/credential_provider';
Expand Down Expand Up @@ -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);
});
};

Expand Down
19 changes: 19 additions & 0 deletions src/plugins/data_source/server/client/configure_client_utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
};
5 changes: 0 additions & 5 deletions src/plugins/data_source/server/legacy/configure_legacy_client.test.mocks.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}));
77 changes: 62 additions & 15 deletions src/plugins/data_source/server/legacy/configure_legacy_client.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -119,10 +130,13 @@ describe('configureLegacyClient', () => {
response: mockResponse,
});
});

authenticationMethodRegistery.getAuthenticationMethod.mockImplementation(() => authMethod);
});

afterEach(() => {
ClientMock.mockReset();
CredentialsMock.mockReset();
jest.resetAllMocks();
});

Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
});
});
});
8 changes: 5 additions & 3 deletions src/plugins/data_source/server/legacy/configure_legacy_client.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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';
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 051087d

Please sign in to comment.