Skip to content

Commit 90aa77a

Browse files
qhanamqhanam
authored
fix: Invalidate cognito identity and re-try (#498)
1 parent 8fc61db commit 90aa77a

File tree

6 files changed

+232
-43
lines changed

6 files changed

+232
-43
lines changed

‎src/dispatch/BasicAuthentication.ts‎

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,21 +27,27 @@ export class BasicAuthentication extends Authentication {
2727
*/
2828
protected AnonymousCognitoCredentialsProvider =
2929
async (): Promise<AwsCredentialIdentity> => {
30-
return this.cognitoIdentityClient
31-
.getId({
32-
IdentityPoolId: this.config.identityPoolId as string
33-
})
34-
.then((getIdResponse) =>
35-
this.cognitoIdentityClient.getOpenIdToken(getIdResponse)
36-
)
37-
.then((getOpenIdTokenResponse) =>
38-
this.stsClient.assumeRoleWithWebIdentity({
39-
RoleArn: this.config.guestRoleArn as string,
40-
RoleSessionName: 'cwr',
41-
WebIdentityToken: getOpenIdTokenResponse.Token
42-
})
43-
)
44-
.then((credentials: AwsCredentialIdentity) => {
30+
let retries = 1;
31+
32+
while (true) {
33+
try {
34+
const getIdResponse =
35+
await this.cognitoIdentityClient.getId({
36+
IdentityPoolId: this.config.identityPoolId as string
37+
});
38+
39+
const getOpenIdTokenResponse =
40+
await this.cognitoIdentityClient.getOpenIdToken(
41+
getIdResponse
42+
);
43+
44+
const credentials =
45+
await this.stsClient.assumeRoleWithWebIdentity({
46+
RoleArn: this.config.guestRoleArn as string,
47+
RoleSessionName: 'cwr',
48+
WebIdentityToken: getOpenIdTokenResponse.Token
49+
});
50+
4551
this.credentials = credentials;
4652
try {
4753
localStorage.setItem(
@@ -51,7 +57,15 @@ export class BasicAuthentication extends Authentication {
5157
} catch (e) {
5258
// Ignore
5359
}
60+
5461
return credentials;
55-
});
62+
} catch (e) {
63+
if (retries) {
64+
retries--;
65+
} else {
66+
throw e;
67+
}
68+
}
69+
}
5670
};
5771
}

‎src/dispatch/CognitoIdentityClient.ts‎

Lines changed: 31 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,11 @@ export class CognitoIdentityClient {
100100
const { response } = await this.fetchRequestHandler.handle(
101101
tokenRequest
102102
);
103-
return (await responseToJson(response)) as OpenIdTokenResponse;
103+
return this.validateOpenIdTokenResponse(
104+
await responseToJson(response)
105+
);
104106
} catch (e) {
107+
localStorage.removeItem(IDENTITY_KEY);
105108
throw new Error(
106109
`CWR: Failed to retrieve Cognito OpenId token: ${e}`
107110
);
@@ -120,39 +123,51 @@ export class CognitoIdentityClient {
120123
const { response } = await this.fetchRequestHandler.handle(
121124
credentialRequest
122125
);
123-
const credentialsResponse = (await responseToJson(
124-
response
125-
)) as CredentialsResponse;
126-
this.validateCredenentialsResponse(credentialsResponse);
127-
const Credentials = credentialsResponse.Credentials;
128126
const { AccessKeyId, Expiration, SecretKey, SessionToken } =
129-
Credentials;
127+
this.validateCredenentialsResponse(
128+
await responseToJson(response)
129+
);
130130
return {
131131
accessKeyId: AccessKeyId as string,
132132
secretAccessKey: SecretKey as string,
133133
sessionToken: SessionToken as string,
134134
expiration: new Date(Expiration * 1000)
135135
};
136136
} catch (e) {
137+
localStorage.removeItem(IDENTITY_KEY);
137138
throw new Error(
138139
`CWR: Failed to retrieve credentials for Cognito identity: ${e}`
139140
);
140141
}
141142
};
142143

143-
private validateCredenentialsResponse = (cr: any) => {
144-
if (
145-
cr &&
146-
cr.__type &&
147-
(cr.__type === 'ResourceNotFoundException' ||
148-
cr.__type === 'ValidationException')
149-
) {
144+
private validateOpenIdTokenResponse = (r: any): OpenIdTokenResponse => {
145+
if ('IdentityId' in r && 'Token' in r) {
146+
return r as OpenIdTokenResponse;
147+
} else if (r && '__type' in r && 'message' in r) {
150148
// The request may have failed because of ValidationException or
151149
// ResourceNotFoundException, which means the identity Id is bad. In
152150
// any case, we invalidate the identity Id so the entire process can
153151
// be re-tried.
154-
localStorage.removeItem(IDENTITY_KEY);
155-
throw new Error(`${cr.__type}: ${cr.message}`);
152+
throw new Error(`${r.__type}: ${r.message}`);
153+
} else {
154+
// We don't recognize ths response format.
155+
throw new Error('Unknown OpenIdToken response');
156+
}
157+
};
158+
159+
private validateCredenentialsResponse = (r: any): CognitoCredentials => {
160+
if ('IdentityId' in r && 'Credentials' in r) {
161+
return (r as CredentialsResponse).Credentials;
162+
} else if (r && '__type' in r && 'message' in r) {
163+
// The request may have failed because of ValidationException or
164+
// ResourceNotFoundException, which means the identity Id is bad. In
165+
// any case, we invalidate the identity Id so the entire process can
166+
// be re-tried.
167+
throw new Error(`${r.__type}: ${r.message}`);
168+
} else {
169+
// We don't recognize ths response format.
170+
throw new Error('Unknown Credentials response');
156171
}
157172
};
158173

‎src/dispatch/EnhancedAuthentication.ts‎

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,20 @@ export class EnhancedAuthentication extends Authentication {
1717
*/
1818
protected AnonymousCognitoCredentialsProvider =
1919
async (): Promise<AwsCredentialIdentity> => {
20-
return this.cognitoIdentityClient
21-
.getId({ IdentityPoolId: this.config.identityPoolId as string })
22-
.then((getIdResponse) =>
23-
this.cognitoIdentityClient.getCredentialsForIdentity(
24-
getIdResponse.IdentityId
25-
)
26-
)
27-
.then((credentials: AwsCredentialIdentity) => {
20+
let retries = 1;
21+
22+
while (true) {
23+
try {
24+
const getIdResponse =
25+
await this.cognitoIdentityClient.getId({
26+
IdentityPoolId: this.config.identityPoolId as string
27+
});
28+
29+
const credentials =
30+
await this.cognitoIdentityClient.getCredentialsForIdentity(
31+
getIdResponse.IdentityId
32+
);
33+
2834
this.credentials = credentials;
2935
try {
3036
localStorage.setItem(
@@ -34,7 +40,15 @@ export class EnhancedAuthentication extends Authentication {
3440
} catch (e) {
3541
// Ignore
3642
}
43+
3744
return credentials;
38-
});
45+
} catch (e) {
46+
if (retries) {
47+
retries--;
48+
} else {
49+
throw e;
50+
}
51+
}
52+
}
3953
};
4054
}

‎src/dispatch/__tests__/BasicAuthentication.test.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,4 +407,31 @@ describe('BasicAuthentication tests', () => {
407407
storageExpiration.getTime()
408408
);
409409
});
410+
411+
test('when mockGetIdToken fails then retry', async () => {
412+
const e: Error = new Error('mockGetId error');
413+
mockGetIdToken.mockImplementationOnce(() => {
414+
throw e;
415+
});
416+
// Init
417+
const auth = new BasicAuthentication({
418+
...DEFAULT_CONFIG,
419+
...{
420+
identityPoolId: IDENTITY_POOL_ID,
421+
guestRoleArn: GUEST_ROLE_ARN
422+
}
423+
});
424+
425+
// Run
426+
const credentials = await auth.ChainAnonymousCredentialsProvider();
427+
428+
// Assert
429+
expect(credentials).toEqual(
430+
expect.objectContaining({
431+
accessKeyId: 'x',
432+
secretAccessKey: 'y',
433+
sessionToken: 'z'
434+
})
435+
);
436+
});
410437
});

‎src/dispatch/__tests__/CognitoIdentityClient.test.ts‎

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -251,16 +251,108 @@ describe('CognitoIdentityClient tests', () => {
251251
).rejects.toEqual(expected);
252252
});
253253

254-
test('when getCredentialsForIdentity returns a ResourceNotFoundException then identity id is removed from localStorage ', async () => {
254+
test('when getCredentialsForIdentity returns bad response then an error is thrown', async () => {
255+
fetchHandler.mockResolvedValueOnce({
256+
response: {
257+
body: getReadableStream('{}')
258+
}
259+
});
260+
const expected: Error = new Error(
261+
`CWR: Failed to retrieve credentials for Cognito identity: Error: Unknown Credentials response`
262+
);
263+
264+
// Init
265+
const client: CognitoIdentityClient = new CognitoIdentityClient({
266+
fetchRequestHandler: new FetchHttpHandler(),
267+
region: Utils.AWS_RUM_REGION
268+
});
269+
270+
// Assert
271+
await expect(
272+
client.getCredentialsForIdentity('my-fake-identity-id')
273+
).rejects.toEqual(expected);
274+
});
275+
276+
test('when getCredentialsForIdentity returns bad response then identity id is removed from localStorage ', async () => {
255277
localStorage.setItem(IDENTITY_KEY, 'my-fake-identity-id');
256278

279+
fetchHandler.mockResolvedValueOnce({
280+
response: {
281+
body: getReadableStream('not-json')
282+
}
283+
});
284+
285+
// Init
286+
const client: CognitoIdentityClient = new CognitoIdentityClient({
287+
fetchRequestHandler: new FetchHttpHandler(),
288+
region: Utils.AWS_RUM_REGION
289+
});
290+
291+
// Run
292+
try {
293+
await client.getCredentialsForIdentity('my-fake-identity-id');
294+
} catch (e) {
295+
// Ignore
296+
}
297+
298+
// Assert
299+
expect(localStorage.getItem(IDENTITY_KEY)).toBe(null);
300+
});
301+
302+
test('when getOpenIdToken returns a ResourceNotFoundException then an error is thrown', async () => {
257303
fetchHandler.mockResolvedValueOnce({
258304
response: {
259305
body: getReadableStream(
260306
'{"__type": "ResourceNotFoundException", "message": ""}'
261307
)
262308
}
263309
});
310+
const expected: Error = new Error(
311+
`CWR: Failed to retrieve Cognito OpenId token: Error: ResourceNotFoundException: `
312+
);
313+
314+
// Init
315+
const client: CognitoIdentityClient = new CognitoIdentityClient({
316+
fetchRequestHandler: new FetchHttpHandler(),
317+
region: Utils.AWS_RUM_REGION
318+
});
319+
320+
// Assert
321+
await expect(
322+
client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' })
323+
).rejects.toEqual(expected);
324+
});
325+
326+
test('when getOpenIdToken returns a bad response then an error is thrown', async () => {
327+
fetchHandler.mockResolvedValueOnce({
328+
response: {
329+
body: getReadableStream('{}')
330+
}
331+
});
332+
const expected: Error = new Error(
333+
`CWR: Failed to retrieve Cognito OpenId token: Error: Unknown OpenIdToken response`
334+
);
335+
336+
// Init
337+
const client: CognitoIdentityClient = new CognitoIdentityClient({
338+
fetchRequestHandler: new FetchHttpHandler(),
339+
region: Utils.AWS_RUM_REGION
340+
});
341+
342+
// Assert
343+
await expect(
344+
client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' })
345+
).rejects.toEqual(expected);
346+
});
347+
348+
test('when getOpenIdToken returns a bad response then identity id is removed from localStorage ', async () => {
349+
localStorage.setItem(IDENTITY_KEY, 'my-fake-identity-id');
350+
351+
fetchHandler.mockResolvedValueOnce({
352+
response: {
353+
body: getReadableStream('not-json')
354+
}
355+
});
264356

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

271363
// Run
272364
try {
273-
await client.getCredentialsForIdentity('my-fake-identity-id');
365+
await client.getOpenIdToken({ IdentityId: 'my-fake-identity-id' });
274366
} catch (e) {
275367
// Ignore
276368
}

‎src/dispatch/__tests__/EnhancedAuthentication.test.ts‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -380,4 +380,31 @@ describe('EnhancedAuthentication tests', () => {
380380
storageExpiration.getTime()
381381
);
382382
});
383+
384+
test('when getCredentialsForIdentity fails then retry', async () => {
385+
// Init
386+
mockGetId.mockImplementationOnce(() => {
387+
throw new Error('mockGetId error');
388+
});
389+
390+
const auth = new EnhancedAuthentication({
391+
...DEFAULT_CONFIG,
392+
...{
393+
identityPoolId: IDENTITY_POOL_ID,
394+
guestRoleArn: GUEST_ROLE_ARN
395+
}
396+
});
397+
398+
// Run
399+
const credentials = await auth.ChainAnonymousCredentialsProvider();
400+
401+
// Assert
402+
expect(credentials).toEqual(
403+
expect.objectContaining({
404+
accessKeyId: 'x',
405+
secretAccessKey: 'y',
406+
sessionToken: 'z'
407+
})
408+
);
409+
});
383410
});

0 commit comments

Comments
 (0)