fix: fix support for custom Cognito redirect path (#87)

In case all paths in a web domain are not handled by the `cognito-at-edge`-based edge function, it is necessary to specify the path to which Cognito will redirect while authenticating. This was already available with the `parseAuthPath` parameter, but it was used inconsistently in the `Authenticator.handle` function, resulting in `Unauthorized` errors from the Cognito API. It is now used consistently when providing a `redirect_uri` to the API. The readme was also updated to provide more guidance about using this option. --------- Co-authored-by: Jean de Kernier <[email protected]>
awslabs · Mar 1, 2024 · 2ddc133 · 2ddc133
1 parent 211518a
commit 2ddc133
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 13 deletions.
diff --git a/README.md b/README.md
@@ -60,6 +60,7 @@ For an explanation of the interactions between CloudFront, Cognito and Lambda@Ed
   * `disableCookieDomain` *boolean* (Optional) Sets domain attribute in cookies, defaults to false (eg: `false`)
   * `httpOnly` *boolean* (Optional) Forbids JavaScript from accessing the cookies, defaults to false (eg: `false`). Note, if this is set to `true`, the cookies will not be accessible to Amplify auth if you are using it client side.
   * `sameSite` *Strict | Lax | None* (Optional) Allows you to declare if your cookie should be restricted to a first-party or same-site context (eg: `SameSite=None`).
+  * `parseAuthPath` *string* (Optional) URI path used as redirect target after successful Cognito authentication (eg: `/oauth2/idpresponse`), defaults to the web domain root. Needs to be a path that is handled by the library. When using this parameter, you should also provide a value for `cookiePath` to ensure your cookies are available for the right paths.
   * `cookiePath` *string* (Optional) Sets Path attribute in cookies
   * `cookieDomain` *string* (Optional) Sets the domain name used for the token cookies
   * `cookieSettingsOverrides` *object* (Optional) Cookie settings overrides for different token cookies -- idToken, accessToken and refreshToken
@@ -73,7 +74,6 @@ For an explanation of the interactions between CloudFront, Cognito and Lambda@Ed
   * `logoutConfiguration` *object* (Optional) Enables logout functionality
     * `logoutUri` *string* URI path, which when matched with request, logs user out by revoking tokens and clearing cookies
     * `logoutRedirectUri` *string* The URI to which the user is redirected to after logging them out
-  * `parseAuthPath` *string* (Optional) URI path to use for the parse auth handler, when the library is used in an authentication gateway setup
   * `csrfProtection` *object* (Optional) Enables CSRF protection
     * `nonceSigningSecret` *string* Secret used for signing nonce cookies
   * `logLevel` *string* (Optional) Logging level. Default: `'silent'`. One of `'fatal'`, `'error'`, `'warn'`, `'info'`, `'debug'`, `'trace'` or `'silent'`.

diff --git a/__tests__/index.test.ts b/__tests__/index.test.ts
@@ -714,6 +714,31 @@ describe('handle', () => {
       });
   });
 
+  test('should fetch and set token if code is present (custom redirect)', () => {
+    const authenticatorWithCustomRedirect : any = new Authenticator({
+      region: 'us-east-1',
+      userPoolId: 'us-east-1_abcdef123',
+      userPoolAppId: '123456789qwertyuiop987abcd',
+      userPoolDomain: 'my-cognito-domain.auth.us-east-1.amazoncognito.com',
+      parseAuthPath: '/custom/login/path',
+    });
+    jest.spyOn(authenticatorWithCustomRedirect._jwtVerifier, 'verify');
+    jest.spyOn(authenticatorWithCustomRedirect, '_fetchTokensFromCode');
+    jest.spyOn(authenticatorWithCustomRedirect, '_getRedirectResponse');
+    authenticatorWithCustomRedirect._jwtVerifier.verify.mockImplementationOnce(async () => { throw new Error(); });
+    authenticatorWithCustomRedirect._fetchTokensFromCode.mockResolvedValueOnce(tokenData);
+    authenticatorWithCustomRedirect._getRedirectResponse.mockReturnValueOnce({ response: 'toto' });
+    const request = getCloudfrontRequest();
+    request.Records[0].cf.request.querystring = 'code=54fe5f4e&state=/lol';
+    return expect(authenticatorWithCustomRedirect.handle(request)).resolves.toEqual({ response: 'toto' })
+      .then(() => {
+        expect(authenticatorWithCustomRedirect._jwtVerifier.verify).toHaveBeenCalled();
+        expect(authenticatorWithCustomRedirect._fetchTokensFromCode).toHaveBeenCalledWith('https://d111111abcdef8.cloudfront.net/custom/login/path', '54fe5f4e');
+        expect(authenticatorWithCustomRedirect._getRedirectResponse).toHaveBeenCalledWith(tokenData, 'd111111abcdef8.cloudfront.net', '/lol');
+      });
+  });
+
+
   test('should fetch and set token if code is present and when csrfProtection is enabled', () => {
     authenticator._jwtVerifier.verify.mockImplementationOnce(async () => { throw new Error(); });
     authenticator._fetchTokensFromCode.mockResolvedValueOnce(tokenData);
@@ -760,6 +785,40 @@ describe('handle', () => {
       });
   });
 
+  test('should redirect to auth domain if unauthenticated and no code (custom redirect)', () => {
+    const authenticatorWithCustomRedirect : any = new Authenticator({
+      region: 'us-east-1',
+      userPoolId: 'us-east-1_abcdef123',
+      userPoolAppId: '123456789qwertyuiop987abcd',
+      userPoolDomain: 'my-cognito-domain.auth.us-east-1.amazoncognito.com',
+      parseAuthPath: '/custom/login/path',
+    });
+    jest.spyOn(authenticatorWithCustomRedirect._jwtVerifier, 'verify');
+    authenticator._jwtVerifier.verify.mockImplementationOnce(async () => { throw new Error();});
+    return expect(authenticatorWithCustomRedirect.handle(getCloudfrontRequest())).resolves.toEqual(
+      {
+        status: '302',
+        headers: {
+          'location': [{
+            key: 'Location',
+            value: 'https://my-cognito-domain.auth.us-east-1.amazoncognito.com/authorize?redirect_uri=https://d111111abcdef8.cloudfront.net/custom/login/path&response_type=code&client_id=123456789qwertyuiop987abcd&state=/lol%3F%3Fparam%3D1',
+          }],
+          'cache-control': [{
+            key: 'Cache-Control',
+            value: 'no-cache, no-store, max-age=0, must-revalidate',
+          }],
+          'pragma': [{
+            key: 'Pragma',
+            value: 'no-cache',
+          }],
+        },
+      },
+    )
+      .then(() => {
+        expect(authenticatorWithCustomRedirect._jwtVerifier.verify).toHaveBeenCalled();
+      });
+  });
+
   test('should redirect to auth domain and clear csrf cookies if unauthenticated and no code', async () => {
     authenticator._jwtVerifier.verify.mockImplementationOnce(async () => { throw new Error(); });
     authenticator._csrfProtection = {
@@ -795,6 +854,21 @@ describe('handle', () => {
     expect(cookies.find(c => c.match(`.${PKCE_COOKIE_NAME_SUFFIX}=`))).toBeDefined();
   });
 
+  test('should redirect to auth domain with custom return redirect if unauthenticated', async () => {
+    const authenticatorWithCustomRedirect : any = new Authenticator({
+      region: 'us-east-1',
+      userPoolId: 'us-east-1_abcdef123',
+      userPoolAppId: '123456789qwertyuiop987abcd',
+      userPoolDomain: 'my-cognito-domain.auth.us-east-1.amazoncognito.com',
+      parseAuthPath: '/custom/login/path',
+    });
+    jest.spyOn(authenticatorWithCustomRedirect._jwtVerifier, 'verify');
+    authenticatorWithCustomRedirect._jwtVerifier.verify.mockImplementationOnce(async () => { throw new Error(); });
+    const response = await authenticatorWithCustomRedirect.handle(getCloudfrontRequest());
+    const url = new URL(response.headers['location'][0].value);
+    expect(url.searchParams.get('redirect_uri')).toEqual('https://d111111abcdef8.cloudfront.net/custom/login/path');
+  });
+
   test('should revoke tokens and clear cookies if logoutConfiguration is set', () => {
     authenticator._logoutConfiguration = { logoutUri: '/logout' };
     authenticator._getTokensFromCookie.mockReturnValueOnce({ refreshToken: tokenData.refresh_token });

diff --git a/src/index.ts b/src/index.ts
@@ -526,25 +526,19 @@ export class Authenticator {
    * @return {CloudFrontResultResponse} Redirect response.
    */
   _getRedirectToCognitoUserPoolResponse(request: CloudFrontRequest, redirectURI: string): CloudFrontResultResponse {
-    const cfDomain = request.headers.host[0].value;
     let redirectPath = request.uri;
     if (request.querystring && request.querystring !== '') {
       redirectPath += encodeURIComponent('?' + request.querystring);
     }
 
-    let oauthRedirectUri = redirectURI;
-    if (this._parseAuthPath) {
-      oauthRedirectUri = `https://${cfDomain}/${this._parseAuthPath}`;
-    }
-
     let csrfTokens: CSRFTokens = {};
     let state: string | undefined = redirectPath;
     if (this._csrfProtection) {
       csrfTokens = generateCSRFTokens(redirectURI, this._csrfProtection.nonceSigningSecret);
       state = csrfTokens.state;
     }
 
-    const userPoolUrl = `https://${this._userPoolDomain}/authorize?redirect_uri=${oauthRedirectUri}&response_type=code&client_id=${this._userPoolAppId}&state=${state}`;
+    const userPoolUrl = `https://${this._userPoolDomain}/authorize?redirect_uri=${redirectURI}&response_type=code&client_id=${this._userPoolAppId}&state=${state}`;
 
     this._logger.debug(`Redirecting user to Cognito User Pool URL ${userPoolUrl}`);
 
@@ -601,9 +595,8 @@ export class Authenticator {
     this._logger.debug({ msg: 'Handling Lambda@Edge event', event });
 
     const { request } = event.Records[0].cf;
-    const requestParams = parse(request.querystring);
     const cfDomain = request.headers.host[0].value;
-    const redirectURI = `https://${cfDomain}`;
+    const redirectURI = this._parseAuthPath ? `https://${cfDomain}/${this._parseAuthPath}` : `https://${cfDomain}`;
 
     try {
       const tokens = this._getTokensFromCookie(request.headers.cookie);
@@ -631,10 +624,12 @@ export class Authenticator {
       }
     } catch (err) {
       if (this._logoutConfiguration && request.uri.startsWith(this._logoutConfiguration.logoutUri)) {
-        this._logger.info({ msg: 'Clearing cookies', path: redirectURI });
+        this._logger.info({ msg: 'Clearing cookies', path: cfDomain });
         return this._clearCookies(event);
       }
       this._logger.debug("User isn't authenticated: %s", err);
+
+      const requestParams = parse(request.querystring);
       if (requestParams.code) {
         return this._fetchTokensFromCode(redirectURI, requestParams.code as string)
           .then(tokens => this._getRedirectResponse(tokens, cfDomain, this._getRedirectUriFromState(requestParams.state as string)));
@@ -678,7 +673,9 @@ export class Authenticator {
       };
     } catch (err) {
       this._logger.debug("User isn't authenticated: %s", err);
-      return this._getRedirectToCognitoUserPoolResponse(request, redirectURI);
+      return this._getRedirectToCognitoUserPoolResponse(
+        request, this._parseAuthPath ? `https://${cfDomain}/${this._parseAuthPath}` : redirectURI,
+      );
     }
   }
 
@@ -756,7 +753,9 @@ export class Authenticator {
       return this._getRedirectResponse(tokens, cfDomain, redirectURI);
     } catch (err) {
       this._logger.debug("User isn't authenticated: %s", err);
-      return this._getRedirectToCognitoUserPoolResponse(request, redirectURI);
+      return this._getRedirectToCognitoUserPoolResponse(
+        request, this._parseAuthPath ? `https://${cfDomain}/${this._parseAuthPath}` : redirectURI,
+      );
     }
   }