Skip to content

Commit

Permalink
fix: Invalidate cognito identity and re-try (#498)
Browse files Browse the repository at this point in the history
  • Loading branch information
@qhanam
qhanam committed Jan 30, 2024
1 parent 26238a5 commit afe87e8
Show file tree
Hide file tree
Showing 6 changed files with 232 additions and 43 deletions.
46 changes: 30 additions & 16 deletions src/dispatch/BasicAuthentication.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,21 +27,27 @@ export class BasicAuthentication extends Authentication {
*/
protected AnonymousCognitoCredentialsProvider =
async (): Promise<AwsCredentialIdentity> => {
return this.cognitoIdentityClient
.getId({
IdentityPoolId: this.config.identityPoolId as string
})
.then((getIdResponse) =>
this.cognitoIdentityClient.getOpenIdToken(getIdResponse)
)
.then((getOpenIdTokenResponse) =>
this.stsClient.assumeRoleWithWebIdentity({
RoleArn: this.config.guestRoleArn as string,
RoleSessionName: 'cwr',
WebIdentityToken: getOpenIdTokenResponse.Token
})
)
.then((credentials: AwsCredentialIdentity) => {
let retries = 1;

while (true) {
try {
const getIdResponse =
await this.cognitoIdentityClient.getId({
IdentityPoolId: this.config.identityPoolId as string
});

const getOpenIdTokenResponse =
await this.cognitoIdentityClient.getOpenIdToken(
getIdResponse
);

const credentials =
await this.stsClient.assumeRoleWithWebIdentity({
RoleArn: this.config.guestRoleArn as string,
RoleSessionName: 'cwr',
WebIdentityToken: getOpenIdTokenResponse.Token
});

this.credentials = credentials;
try {
localStorage.setItem(
Expand All @@ -51,7 +57,15 @@ export class BasicAuthentication extends Authentication {
} catch (e) {
// Ignore
}

return credentials;
});
} catch (e) {
if (retries) {
retries--;
} else {
throw e;
}
}
}
};
}
47 changes: 31 additions & 16 deletions src/dispatch/CognitoIdentityClient.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,8 +100,11 @@ export class CognitoIdentityClient {
const { response } = await this.fetchRequestHandler.handle(
tokenRequest
);
return (await responseToJson(response)) as OpenIdTokenResponse;
return this.validateOpenIdTokenResponse(
await responseToJson(response)
);
} catch (e) {
localStorage.removeItem(IDENTITY_KEY);
throw new Error(
`CWR: Failed to retrieve Cognito OpenId token: ${e}`
);
Expand All @@ -120,39 +123,51 @@ export class CognitoIdentityClient {
const { response } = await this.fetchRequestHandler.handle(
credentialRequest
);
const credentialsResponse = (await responseToJson(
response
)) as CredentialsResponse;
this.validateCredenentialsResponse(credentialsResponse);
const Credentials = credentialsResponse.Credentials;
const { AccessKeyId, Expiration, SecretKey, SessionToken } =
Credentials;
this.validateCredenentialsResponse(
await responseToJson(response)
);
return {
accessKeyId: AccessKeyId as string,
secretAccessKey: SecretKey as string,
sessionToken: SessionToken as string,
expiration: new Date(Expiration * 1000)
};
} catch (e) {
localStorage.removeItem(IDENTITY_KEY);
throw new Error(
`CWR: Failed to retrieve credentials for Cognito identity: ${e}`
);
}
};

private validateCredenentialsResponse = (cr: any) => {
if (
cr &&
cr.__type &&
(cr.__type === 'ResourceNotFoundException' ||
cr.__type === 'ValidationException')
) {
private validateOpenIdTokenResponse = (r: any): OpenIdTokenResponse => {
if ('IdentityId' in r && 'Token' in r) {
return r as OpenIdTokenResponse;
} else if (r && '__type' in r && 'message' in r) {
// The request may have failed because of ValidationException or
// ResourceNotFoundException, which means the identity Id is bad. In
// any case, we invalidate the identity Id so the entire process can
// be re-tried.
localStorage.removeItem(IDENTITY_KEY);
throw new Error(`${cr.__type}: ${cr.message}`);
throw new Error(`${r.__type}: ${r.message}`);
} else {
// We don't recognize ths response format.
throw new Error('Unknown OpenIdToken response');
}
};

private validateCredenentialsResponse = (r: any): CognitoCredentials => {
if ('IdentityId' in r && 'Credentials' in r) {
return (r as CredentialsResponse).Credentials;
} else if (r && '__type' in r && 'message' in r) {
// The request may have failed because of ValidationException or
// ResourceNotFoundException, which means the identity Id is bad. In
// any case, we invalidate the identity Id so the entire process can
// be re-tried.
throw new Error(`${r.__type}: ${r.message}`);
} else {
// We don't recognize ths response format.
throw new Error('Unknown Credentials response');
}
};

Expand Down
32 changes: 23 additions & 9 deletions src/dispatch/EnhancedAuthentication.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,20 @@ export class EnhancedAuthentication extends Authentication {
*/
protected AnonymousCognitoCredentialsProvider =
async (): Promise<AwsCredentialIdentity> => {
return this.cognitoIdentityClient
.getId({ IdentityPoolId: this.config.identityPoolId as string })
.then((getIdResponse) =>
this.cognitoIdentityClient.getCredentialsForIdentity(
getIdResponse.IdentityId
)
)
.then((credentials: AwsCredentialIdentity) => {
let retries = 1;

while (true) {
try {
const getIdResponse =
await this.cognitoIdentityClient.getId({
IdentityPoolId: this.config.identityPoolId as string
});

const credentials =
await this.cognitoIdentityClient.getCredentialsForIdentity(
getIdResponse.IdentityId
);

this.credentials = credentials;
try {
localStorage.setItem(
Expand All @@ -34,7 +40,15 @@ export class EnhancedAuthentication extends Authentication {
} catch (e) {
// Ignore
}

return credentials;
});
} catch (e) {
if (retries) {
retries--;
} else {
throw e;
}
}
}
};
}
27 changes: 27 additions & 0 deletions src/dispatch/__tests__/BasicAuthentication.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -407,4 +407,31 @@ describe('BasicAuthentication tests', () => {
storageExpiration.getTime()
);
});

test('when mockGetIdToken fails then retry', async () => {
const e: Error = new Error('mockGetId error');
mockGetIdToken.mockImplementationOnce(() => {
throw e;
});
// Init
const auth = new BasicAuthentication({
...DEFAULT_CONFIG,
...{
identityPoolId: IDENTITY_POOL_ID,
guestRoleArn: GUEST_ROLE_ARN
}
});

// Run
const credentials = await auth.ChainAnonymousCredentialsProvider();

// Assert
expect(credentials).toEqual(
expect.objectContaining({
accessKeyId: 'x',
secretAccessKey: 'y',
sessionToken: 'z'
})
);
});
});
96 changes: 94 additions & 2 deletions src/dispatch/__tests__/CognitoIdentityClient.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -251,16 +251,108 @@ describe('CognitoIdentityClient tests', () => {
).rejects.toEqual(expected);
});

test('when getCredentialsForIdentity returns a ResourceNotFoundException then identity id is removed from localStorage ', async () => {
test('when getCredentialsForIdentity returns bad response then an error is thrown', async () => {
fetchHandler.mockResolvedValueOnce({
response: {
body: getReadableStream('{}')
}
});
const expected: Error = new Error(
`CWR: Failed to retrieve credentials for Cognito identity: Error: Unknown Credentials response`
);

// Init
const client: CognitoIdentityClient = new CognitoIdentityClient({
fetchRequestHandler: new FetchHttpHandler(),
region: Utils.AWS_RUM_REGION
});

// Assert
await expect(
client.getCredentialsForIdentity('my-fake-identity-id')
).rejects.toEqual(expected);
});

test('when getCredentialsForIdentity returns bad response then identity id is removed from localStorage ', async () => {
localStorage.setItem(IDENTITY_KEY, 'my-fake-identity-id');

fetchHandler.mockResolvedValueOnce({
response: {
body: getReadableStream('not-json')
}
});

// Init
const client: CognitoIdentityClient = new CognitoIdentityClient({
fetchRequestHandler: new FetchHttpHandler(),
region: Utils.AWS_RUM_REGION
});

// Run
try {
await client.getCredentialsForIdentity('my-fake-identity-id');
} catch (e) {
// Ignore
}

// Assert
expect(localStorage.getItem(IDENTITY_KEY)).toBe(null);
});

test('when getOpenIdToken returns a ResourceNotFoundException then an error is thrown', async () => {
fetchHandler.mockResolvedValueOnce({
response: {
body: getReadableStream(
'{"__type": "ResourceNotFoundException", "message": ""}'
)
}
});
const expected: Error = new Error(
`CWR: Failed to retrieve Cognito OpenId token: Error: ResourceNotFoundException: `
);

// Init
const client: CognitoIdentityClient = new CognitoIdentityClient({
fetchRequestHandler: new FetchHttpHandler(),
region: Utils.AWS_RUM_REGION
});

// Assert
await expect(
client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' })
).rejects.toEqual(expected);
});

test('when getOpenIdToken returns a bad response then an error is thrown', async () => {
fetchHandler.mockResolvedValueOnce({
response: {
body: getReadableStream('{}')
}
});
const expected: Error = new Error(
`CWR: Failed to retrieve Cognito OpenId token: Error: Unknown OpenIdToken response`
);

// Init
const client: CognitoIdentityClient = new CognitoIdentityClient({
fetchRequestHandler: new FetchHttpHandler(),
region: Utils.AWS_RUM_REGION
});

// Assert
await expect(
client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' })
).rejects.toEqual(expected);
});

test('when getOpenIdToken returns a bad response then identity id is removed from localStorage ', async () => {
localStorage.setItem(IDENTITY_KEY, 'my-fake-identity-id');

fetchHandler.mockResolvedValueOnce({
response: {
body: getReadableStream('not-json')
}
});

// Init
const client: CognitoIdentityClient = new CognitoIdentityClient({
Expand All @@ -270,7 +362,7 @@ describe('CognitoIdentityClient tests', () => {

// Run
try {
await client.getCredentialsForIdentity('my-fake-identity-id');
await client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' });
} catch (e) {
// Ignore
}
Expand Down
27 changes: 27 additions & 0 deletions src/dispatch/__tests__/EnhancedAuthentication.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -380,4 +380,31 @@ describe('EnhancedAuthentication tests', () => {
storageExpiration.getTime()
);
});

test('when getCredentialsForIdentity fails then retry', async () => {
// Init
mockGetId.mockImplementationOnce(() => {
throw new Error('mockGetId error');
});

const auth = new EnhancedAuthentication({
...DEFAULT_CONFIG,
...{
identityPoolId: IDENTITY_POOL_ID,
guestRoleArn: GUEST_ROLE_ARN
}
});

// Run
const credentials = await auth.ChainAnonymousCredentialsProvider();

// Assert
expect(credentials).toEqual(
expect.objectContaining({
accessKeyId: 'x',
secretAccessKey: 'y',
sessionToken: 'z'
})
);
});
});

0 comments on commit afe87e8

Please sign in to comment.