Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport 2.x] Adds a session token to AWS Credentials #6112

Merged
merged 1 commit into from
Mar 11, 2024
Merged

[Backport 2.x] Adds a session token to AWS Credentials #6112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/plugins/data_source/server/client/configure_client.test.mocks.ts
View file
Open in desktop
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
Open in desktop
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
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 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);

Check warning on line 208 in src/plugins/data_source/server/client/configure_client.ts

View check run for this annotation

Codecov / codecov/patch

src/plugins/data_source/server/client/configure_client.ts#L208 

Added line #L208 was not covered by tests
});
};

Expand Down
19 changes: 19 additions & 0 deletions src/plugins/data_source/server/client/configure_client_utils.ts
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading