Skip to content

Commit

Permalink
feat(auth): support custom identity pool endpoint
Browse files Browse the repository at this point in the history
  • Loading branch information
HuiSF committed Oct 3, 2024
1 parent 78ee7df commit 022f28d
Show file tree
Hide file tree
Showing 7 changed files with 173 additions and 57 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,8 @@
// SPDX-License-Identifier: Apache-2.0

import {
GetCredentialsForIdentityOutput,
ResourcesConfig,
getCredentialsForIdentity,
createGetCredentialsForIdentityClient,
sharedInMemoryStorage,
} from '@aws-amplify/core';

Expand All @@ -18,7 +17,7 @@ import { authAPITestParams } from './testUtils/authApiTestParams';

jest.mock('@aws-amplify/core', () => ({
...jest.requireActual('@aws-amplify/core'),
getCredentialsForIdentity: jest.fn(),
createGetCredentialsForIdentityClient: jest.fn(),
}));

jest.mock(
Expand Down Expand Up @@ -63,24 +62,38 @@ const disallowGuestAccessConfig: ResourcesConfig = {
},
};

const credentialsForIdentityIdSpy = getCredentialsForIdentity as jest.Mock;
const mockCreateGetIdentityForIdentityClient = jest.mocked(
createGetCredentialsForIdentityClient,
);

const mockGetCredentialsForIdentity: jest.MockedFunction<
ReturnType<typeof createGetCredentialsForIdentityClient>
> = jest.fn(
async (_config, _params) => authAPITestParams.CredentialsForIdentityIdResult,
);

describe('Guest Credentials', () => {
let cognitoCredentialsProvider: CognitoAWSCredentialsAndIdentityIdProvider;

beforeAll(() => {
mockCreateGetIdentityForIdentityClient.mockReturnValue(
mockGetCredentialsForIdentity,
);
});

describe('Happy Path Cases:', () => {
beforeEach(() => {
cognitoCredentialsProvider =
new CognitoAWSCredentialsAndIdentityIdProvider(
new DefaultIdentityIdStore(sharedInMemoryStorage),
);
credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
return authAPITestParams.CredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockImplementationOnce(async () => {
return authAPITestParams.CredentialsForIdentityIdResult;
});
});
afterEach(() => {
cognitoCredentialsProvider.clearCredentials();
credentialsForIdentityIdSpy?.mockReset();
mockGetCredentialsForIdentity?.mockReset();
});
test('Should call identityIdClient with no logins to obtain guest creds', async () => {
const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
Expand All @@ -92,8 +105,8 @@ describe('Guest Credentials', () => {
.AccessKeyId,
);

expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(credentialsForIdentityIdSpy).toHaveBeenCalledWith(
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledWith(
{ region: 'us-east-1' },
{ IdentityId: 'identity-id-test' },
);
Expand All @@ -107,7 +120,7 @@ describe('Guest Credentials', () => {
authenticated: false,
authConfig: validAuthConfig.Auth!,
});
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);
const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
authenticated: false,
authConfig: validAuthConfig.Auth!,
Expand All @@ -117,7 +130,7 @@ describe('Guest Credentials', () => {
.AccessKeyId,
);
// expecting to be called only once becasue in-memory creds should be returned
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);
});
});

Expand All @@ -127,16 +140,16 @@ describe('Guest Credentials', () => {
new CognitoAWSCredentialsAndIdentityIdProvider(
new DefaultIdentityIdStore(sharedInMemoryStorage),
);
credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
return authAPITestParams.NoAccessKeyCredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockImplementationOnce(async () => {
return authAPITestParams.NoAccessKeyCredentialsForIdentityIdResult;
});
});

afterEach(() => {
cognitoCredentialsProvider.clearCredentials();
});
afterAll(() => {
credentialsForIdentityIdSpy?.mockReset();
mockGetCredentialsForIdentity?.mockReset();
});
test('Should not throw AuthError when allowGuestAccess is false in the config', async () => {
expect(
Expand Down Expand Up @@ -165,13 +178,13 @@ describe('Primary Credentials', () => {
new CognitoAWSCredentialsAndIdentityIdProvider(
new DefaultIdentityIdStore(sharedInMemoryStorage),
);
credentialsForIdentityIdSpy.mockImplementation(async () => {
return authAPITestParams.CredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockImplementation(async () => {
return authAPITestParams.CredentialsForIdentityIdResult;
});
});
afterEach(() => {
cognitoCredentialsProvider.clearCredentials();
credentialsForIdentityIdSpy?.mockReset();
mockGetCredentialsForIdentity?.mockReset();
});
test('Should call identityIdClient with the logins map to obtain primary creds', async () => {
const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
Expand All @@ -184,21 +197,21 @@ describe('Primary Credentials', () => {
.AccessKeyId,
);

expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);
});
test('in-memory primary creds are returned if not expired and not past TTL', async () => {
await cognitoCredentialsProvider.getCredentialsAndIdentityId({
authenticated: true,
authConfig: validAuthConfig.Auth!,
tokens: authAPITestParams.ValidAuthTokens,
});
expect(credentialsForIdentityIdSpy).toHaveBeenCalledWith(
expect(mockGetCredentialsForIdentity).toHaveBeenCalledWith(
{
region: authAPITestParams.CredentialsClientRequest.region,
},
authAPITestParams.CredentialsClientRequest.withValidAuthToken,
);
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);

const res = await cognitoCredentialsProvider.getCredentialsAndIdentityId({
authenticated: true,
Expand All @@ -210,34 +223,34 @@ describe('Primary Credentials', () => {
.AccessKeyId,
);
// expecting to be called only once becasue in-memory creds should be returned
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);
});
test('Should get new credentials when tokens have changed', async () => {
await cognitoCredentialsProvider.getCredentialsAndIdentityId({
authenticated: true,
authConfig: validAuthConfig.Auth!,
tokens: authAPITestParams.ValidAuthTokens,
});
expect(credentialsForIdentityIdSpy).toHaveBeenCalledWith(
expect(mockGetCredentialsForIdentity).toHaveBeenCalledWith(
{
region: authAPITestParams.CredentialsClientRequest.region,
},
authAPITestParams.CredentialsClientRequest.withValidAuthToken,
);
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(1);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(1);

await cognitoCredentialsProvider.getCredentialsAndIdentityId({
authenticated: true,
authConfig: validAuthConfig.Auth!,
tokens: authAPITestParams.NewValidAuthTokens,
});
expect(credentialsForIdentityIdSpy).toHaveBeenCalledWith(
expect(mockGetCredentialsForIdentity).toHaveBeenCalledWith(
{
region: authAPITestParams.CredentialsClientRequest.region,
},
authAPITestParams.CredentialsClientRequest.withNewValidAuthToken,
);
expect(credentialsForIdentityIdSpy).toHaveBeenCalledTimes(2);
expect(mockGetCredentialsForIdentity).toHaveBeenCalledTimes(2);
});
});
describe('Error Path Cases:', () => {
Expand All @@ -251,11 +264,11 @@ describe('Primary Credentials', () => {
cognitoCredentialsProvider.clearCredentials();
});
afterAll(() => {
credentialsForIdentityIdSpy?.mockReset();
mockGetCredentialsForIdentity?.mockReset();
});
test('Should throw AuthError if either Credentials, accessKeyId or secretKey is absent in the response', async () => {
credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
return authAPITestParams.NoAccessKeyCredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockImplementationOnce(async () => {
return authAPITestParams.NoAccessKeyCredentialsForIdentityIdResult;
});
expect(
cognitoCredentialsProvider.getCredentialsAndIdentityId({
Expand All @@ -264,9 +277,9 @@ describe('Primary Credentials', () => {
tokens: authAPITestParams.ValidAuthTokens,
}),
).rejects.toThrow(AuthError);
credentialsForIdentityIdSpy.mockClear();
credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
return authAPITestParams.NoCredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockClear();
mockGetCredentialsForIdentity.mockImplementationOnce(async () => {
return authAPITestParams.NoCredentialsForIdentityIdResult;
});
expect(
cognitoCredentialsProvider.getCredentialsAndIdentityId({
Expand All @@ -275,9 +288,9 @@ describe('Primary Credentials', () => {
tokens: authAPITestParams.ValidAuthTokens,
}),
).rejects.toThrow(AuthError);
credentialsForIdentityIdSpy.mockClear();
credentialsForIdentityIdSpy.mockImplementationOnce(async () => {
return authAPITestParams.NoSecretKeyInCredentialsForIdentityIdResult as GetCredentialsForIdentityOutput;
mockGetCredentialsForIdentity.mockClear();
mockGetCredentialsForIdentity.mockImplementationOnce(async () => {
return authAPITestParams.NoSecretKeyInCredentialsForIdentityIdResult;
});
expect(
cognitoCredentialsProvider.getCredentialsAndIdentityId({
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import { AmplifyUrl } from '@aws-amplify/core/internals/utils';
import { cognitoIdentityPoolEndpointResolver } from '@aws-amplify/core';

import { createCognitoIdentityPoolEndpointResolver } from '../../../../src/providers/cognito/factories/createCognitoIdentityPoolEndpointResolver';

jest.mock('@aws-amplify/core');

const mockCognitoIdentityPoolEndpointResolver = jest.mocked(
cognitoIdentityPoolEndpointResolver,
);

describe('createCognitoIdentityPoolEndpointResolver()', () => {
afterEach(() => {
mockCognitoIdentityPoolEndpointResolver.mockClear();
});

describe('creating a resolver with overrideEndpoint as `undefined`', () => {
const resolver = createCognitoIdentityPoolEndpointResolver({
endpointOverride: undefined,
});

it('invokes cognitoUserPoolEndpointResolver with the expected region', () => {
const expectedReturningUrl = {
url: new AmplifyUrl(
'https://cognito-identity.us-west-2.amazonaws.com/',
),
};
mockCognitoIdentityPoolEndpointResolver.mockReturnValueOnce(
expectedReturningUrl,
);

const expectedRegion = 'us-west-2';
const { url } = resolver({ region: expectedRegion });

expect(mockCognitoIdentityPoolEndpointResolver).toHaveBeenCalledWith({
region: expectedRegion,
});
expect(url).toStrictEqual(expectedReturningUrl.url);
});
});

describe('creating a resolver with overrideEndpoint', () => {
const endpointOverride = 'https://cognito-identity.example.com';
const resolver = createCognitoIdentityPoolEndpointResolver({
endpointOverride,
});

it('returns the endpoint override', () => {
const expectedRegion = 'us-west-2';
const { url } = resolver({ region: expectedRegion });
expect(mockCognitoIdentityPoolEndpointResolver).not.toHaveBeenCalled();
expect(url).toStrictEqual(
new AmplifyUrl('https://cognito-identity.example.com'),
);
});
});
});
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

import { Amplify, Identity, ResourcesConfig, getId } from '@aws-amplify/core';
import {
GetIdInput,
GetIdOutput,
} from '@aws-amplify/core/internals/aws-clients/cognitoIdentity';
Amplify,
Identity,
ResourcesConfig,
createGetIdClient,
} from '@aws-amplify/core';
import { CognitoIdentityPoolConfig } from '@aws-amplify/core/internals/utils';

import { DefaultIdentityIdStore } from '../../../src/providers/cognito/credentialsProvider/IdentityIdStore';
Expand All @@ -15,9 +16,8 @@ import { authAPITestParams } from './testUtils/authApiTestParams';

jest.mock('@aws-amplify/core', () => ({
...jest.requireActual('@aws-amplify/core'),
getId: jest.fn(),
createGetIdClient: jest.fn(),
}));
jest.mock('@aws-amplify/core/internals/aws-clients/cognitoIdentity');
jest.mock('../../../src/providers/cognito/credentialsProvider/IdentityIdStore');

const ampConfig: ResourcesConfig = {
Expand All @@ -30,7 +30,7 @@ const ampConfig: ResourcesConfig = {
},
};

const mockGetId = getId as jest.Mock;
const mockCreateGetIdClient = jest.mocked(createGetIdClient);
const mockKeyValueStorage = {
setItem: jest.fn(),
getItem: jest.fn(),
Expand All @@ -43,23 +43,25 @@ describe('Cognito IdentityId Provider Happy Path Cases:', () => {
const _ = new DefaultIdentityIdStore(mockKeyValueStorage);
const mockDefaultIdentityIdStoreInstance =
MockDefaultIdentityIdStore.mock.instances[0];
const mockGetId: jest.MockedFunction<ReturnType<typeof createGetIdClient>> =
jest.fn(async (_config, params) => {
if (params.Logins && Object.keys(params.Logins).length === 0) {
return {
IdentityId: authAPITestParams.GuestIdentityId.id,
$metadata: {},
};
} else {
return {
IdentityId: authAPITestParams.PrimaryIdentityId.id,
$metadata: {},
};
}
});

beforeAll(() => {
jest.spyOn(Amplify, 'getConfig').mockImplementationOnce(() => ampConfig);

mockGetId.mockImplementation(
async (_config: object, params: GetIdInput) => {
if (params.Logins && Object.keys(params.Logins).length === 0) {
return {
IdentityId: authAPITestParams.GuestIdentityId.id,
} as GetIdOutput;
} else {
return {
IdentityId: authAPITestParams.PrimaryIdentityId.id,
} as GetIdOutput;
}
},
);
mockCreateGetIdClient.mockReturnValue(mockGetId);
});

afterEach(() => {
Expand All @@ -80,6 +82,7 @@ describe('Cognito IdentityId Provider Happy Path Cases:', () => {
).toBe(authAPITestParams.GuestIdentityId.id);
expect(mockGetId).toHaveBeenCalledTimes(0);
});

test('Should generate a guest identityId and return it', async () => {
mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
async () => {
Expand All @@ -102,6 +105,7 @@ describe('Cognito IdentityId Provider Happy Path Cases:', () => {
).toBe(authAPITestParams.GuestIdentityId.id);
expect(mockGetId).toHaveBeenCalledTimes(1);
});

test('Should return stored primary identityId', async () => {
mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
async () => {
Expand All @@ -117,6 +121,7 @@ describe('Cognito IdentityId Provider Happy Path Cases:', () => {
).toBe(authAPITestParams.PrimaryIdentityId.id);
expect(mockGetId).toHaveBeenCalledTimes(0);
});

test('Should generate a primary identityId and return it', async () => {
mockDefaultIdentityIdStoreInstance.loadIdentityId.mockImplementationOnce(
async () => {
Expand Down
Loading

0 comments on commit 022f28d

Please sign in to comment.