Skip to content

Commit 1bedc33

Browse files
feat(auth): add custom headers support for Cognito Auth requests
Add optional `headers` function to `LibraryAuthOptions` that allows attaching custom headers to all Auth requests to Cognito, consistent with the existing REST and GraphQL custom headers pattern. Closes #12308 Closes #13197 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 032d6d3 commit 1bedc33

File tree

3 files changed

+45
-3
lines changed

3 files changed

+45
-3
lines changed

packages/auth/__tests__/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler/cognitoUserPoolTransferHandler.test.ts

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
1+
import { Amplify } from '@aws-amplify/core';
12
import { unauthenticatedHandler } from '@aws-amplify/core/internals/aws-client-utils';
23
import { composeTransferHandler } from '@aws-amplify/core/internals/aws-client-utils/composers';
34

45
import { cognitoUserPoolTransferHandler } from '../../../../../../../src/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler';
56

7+
jest.mock('@aws-amplify/core');
68
jest.mock('@aws-amplify/core/internals/aws-client-utils');
79
jest.mock('@aws-amplify/core/internals/aws-client-utils/composers');
810

@@ -16,7 +18,7 @@ describe('cognitoUserPoolTransferHandler', () => {
1618
const _ = cognitoUserPoolTransferHandler;
1719
});
1820

19-
it('adds the disableCacheMiddlewareFactory at module loading', () => {
21+
it('adds the disableCacheMiddlewareFactory at module loading', async () => {
2022
expect(mockComposeTransferHandler).toHaveBeenCalledTimes(1);
2123

2224
const [core, middleware] = mockComposeTransferHandler.mock.calls[0];
@@ -33,11 +35,45 @@ describe('cognitoUserPoolTransferHandler', () => {
3335
headers: {},
3436
};
3537

36-
disableCacheMiddleware(mockRequest);
38+
await disableCacheMiddleware(mockRequest);
3739

3840
expect(mockNext).toHaveBeenCalledWith(mockRequest);
3941
expect(mockRequest.headers).toEqual({
4042
'cache-control': 'no-store',
4143
});
4244
});
45+
46+
it('attaches custom headers from libraryOptions when configured', async () => {
47+
const mockHeaders = jest.fn().mockResolvedValue({
48+
'custom-header': 'custom-value',
49+
});
50+
(Amplify as any).libraryOptions = {
51+
Auth: { headers: mockHeaders },
52+
};
53+
54+
const [, middleware] = mockComposeTransferHandler.mock.calls[0];
55+
const disableCacheMiddlewareFactory = middleware[0] as any;
56+
const middlewareFn = disableCacheMiddlewareFactory()(jest.fn());
57+
const mockRequest = { headers: {} as Record<string, string> };
58+
59+
await middlewareFn(mockRequest);
60+
61+
expect(mockHeaders).toHaveBeenCalled();
62+
expect(mockRequest.headers['custom-header']).toBe('custom-value');
63+
});
64+
65+
it('does not attach custom headers when not configured', async () => {
66+
(Amplify as any).libraryOptions = {};
67+
68+
const [, middleware] = mockComposeTransferHandler.mock.calls[0];
69+
const disableCacheMiddlewareFactory = middleware[0] as any;
70+
const mockNext = jest.fn();
71+
const middlewareFn = disableCacheMiddlewareFactory()(mockNext);
72+
const mockRequest = { headers: {} as Record<string, string> };
73+
74+
await middlewareFn(mockRequest);
75+
76+
expect(mockRequest.headers['custom-header']).toBeUndefined();
77+
expect(mockNext).toHaveBeenCalledWith(mockRequest);
78+
});
4379
});

packages/auth/src/foundation/factories/serviceClients/cognitoIdentityProvider/shared/handler/cognitoUserPoolTransferHandler.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

4+
import { Amplify } from '@aws-amplify/core';
45
import { composeTransferHandler } from '@aws-amplify/core/internals/aws-client-utils/composers';
56
import {
67
HttpRequest,
@@ -18,7 +19,10 @@ const disableCacheMiddlewareFactory: Middleware<
1819
Record<string, unknown>
1920
> = () => (next, _) =>
2021
async function disableCacheMiddleware(request) {
21-
request.headers['cache-control'] = 'no-store';
22+
request.headers = {
23+
'cache-control': 'no-store',
24+
...(await Amplify.libraryOptions?.Auth?.headers?.()),
25+
};
2226

2327
return next(request);
2428
};

packages/core/src/singleton/Auth/types.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

4+
import { Headers } from '../../clients';
45
import { StrictUnion } from '../../types';
56
import { AtLeastOne } from '../types';
67

@@ -56,6 +57,7 @@ export interface AuthSession {
5657
export interface LibraryAuthOptions {
5758
tokenProvider?: TokenProvider;
5859
credentialsProvider?: CredentialsAndIdentityIdProvider;
60+
headers?(): Promise<Headers>;
5961
}
6062

6163
export interface Identity {

0 commit comments

Comments
 (0)