From 2767ac35d10d82e75b2ce4225d3866eff2e7c402 Mon Sep 17 00:00:00 2001
From: sarossilli <sarossilli@gmail.com>
Date: Sat, 28 Sep 2024 22:36:20 +0000
Subject: [PATCH] feat(auth): Add ExtraQueryParams to
 AuthSignInWithRedirectInput

---
 .../cognito/signInWithRedirect.test.ts         | 13 +++++++++++++
 .../cognito/apis/signInWithRedirect.ts         |  5 +++++
 packages/auth/src/types/inputs.ts              | 18 ++++++++++++++++++
 3 files changed, 36 insertions(+)

diff --git a/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts b/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
index 8f91323319f..b943ea184a3 100644
--- a/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
+++ b/packages/auth/__tests__/providers/cognito/signInWithRedirect.test.ts
@@ -176,6 +176,19 @@ describe('signInWithRedirect', () => {
 		expect(mockUrlSafeEncode).toHaveBeenCalledWith(expectedCustomState);
 	});
 
+	it('uses extra query parameters when specified', async () => {
+		const expectedDefaultProvider = 'COGNITO';
+		const exptedParamKey = 'customParam';
+		const expectedParamValue = 'object';
+		await signInWithRedirect({
+			extraQueryParams: { customParam: expectedParamValue },
+		});
+		const [oauthUrl] = mockOpenAuthSession.mock.calls[0];
+		expect(oauthUrl).toStrictEqual(
+			`https://oauth.domain.com/oauth2/authorize?redirect_uri=http%3A%2F%2Flocalhost%3A3000%2F&response_type=code&client_id=userPoolClientId&identity_provider=${expectedDefaultProvider}&scope=phone%20email%20openid%20profile%20aws.cognito.signin.user.admin&state=oauth_state&code_challenge=code_challenge&code_challenge_method=S256&${exptedParamKey}=${expectedParamValue}`,
+		);
+	});
+
 	describe('specifications on Web', () => {
 		describe('side effect', () => {
 			it('attaches oauth listener to the Amplify singleton', async () => {
diff --git a/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts b/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
index cab4f018ee7..fe5e092c274 100644
--- a/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
+++ b/packages/auth/src/providers/cognito/apis/signInWithRedirect.ts
@@ -25,6 +25,7 @@ import {
 } from '../utils/oauth';
 import { createOAuthError } from '../utils/oauth/createOAuthError';
 import { listenForOAuthFlowCancellation } from '../utils/oauth/cancelOAuthFlow';
+import { ExtraQueryParameters } from '../../../types/inputs';
 
 /**
  * Signs in a user with OAuth. Redirects the application to an Identity Provider.
@@ -57,6 +58,7 @@ export async function signInWithRedirect(
 		provider,
 		customState: input?.customState,
 		preferPrivateSession: input?.options?.preferPrivateSession,
+		extraQueryParams: input?.extraQueryParams,
 	});
 }
 
@@ -66,12 +68,14 @@ const oauthSignIn = async ({
 	clientId,
 	customState,
 	preferPrivateSession,
+	extraQueryParams,
 }: {
 	oauthConfig: OAuthConfig;
 	provider: string;
 	clientId: string;
 	customState?: string;
 	preferPrivateSession?: boolean;
+	extraQueryParams?: ExtraQueryParameters;
 }) => {
 	const { domain, redirectSignIn, responseType, scopes } = oauthConfig;
 	const randomState = generateState();
@@ -104,6 +108,7 @@ const oauthSignIn = async ({
 			code_challenge: toCodeChallenge(),
 			code_challenge_method: method,
 		}),
+		...extraQueryParams,
 	})
 		.map(([k, v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`)
 		.join('&');
diff --git a/packages/auth/src/types/inputs.ts b/packages/auth/src/types/inputs.ts
index 6e152cdc1e5..ecc2ab7cc29 100644
--- a/packages/auth/src/types/inputs.ts
+++ b/packages/auth/src/types/inputs.ts
@@ -54,6 +54,23 @@ export interface AuthSignOutInput {
 
 export type AuthProvider = 'Amazon' | 'Apple' | 'Facebook' | 'Google';
 
+export type OAuthUrlReservedKeys =
+	| 'redirect_uri'
+	| 'response_type'
+	| 'client_id'
+	| 'identity_provider'
+	| 'scope'
+	| 'state'
+	| 'responseType';
+
+/**
+ * User-defined custom query params for oAuthUrl.
+ */
+export type ExtraQueryParameters = Omit<
+	Record<string, string>,
+	OAuthUrlReservedKeys
+>;
+
 export interface AuthSignInWithRedirectInput {
 	provider?: AuthProvider | { custom: string };
 	customState?: string;
@@ -69,6 +86,7 @@ export interface AuthSignInWithRedirectInput {
 		 */
 		preferPrivateSession?: boolean;
 	};
+	extraQueryParams?: ExtraQueryParameters;
 }
 
 /**