Add support for custom cognito redirect path

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'`.

__tests__/index.test.ts

@@ -448,7 +448,7 @@
       method: 'POST',
     }));
   });
 
   describe('_clearCookies', () => {
     it('should verify tokens and clear cookies', async () => {
       jest.spyOn(authenticator._jwtVerifier, 'verify');
@@ -463,7 +463,7 @@
       expect(response.headers['set-cookie'].length).toBe(5);
     });
 
     it('should clear cookies even if tokens cannot be verified', async () => {
       jest.spyOn(authenticator._jwtVerifier, 'verify');
       authenticator._jwtVerifier.cacheJwks(jwksData);
       authenticator._jwtVerifier.verify.mockReturnValueOnce(Promise.reject({}));
@@ -480,7 +480,7 @@
 
     it('should clear cookies and redirect to logoutRedirectUri', async () => {
       jest.spyOn(authenticator._jwtVerifier, 'verify');
       authenticator._logoutConfiguration = {
         logoutUri: '/logout',
         logoutRedirectUri: 'https://foobar.com',
       };
@@ -491,7 +491,7 @@
       expect(response).toEqual(expect.objectContaining({ status: '302' }));
       expect(response.headers['location']?.[0]?.value).toEqual('https://foobar.com');
     });
 
     it('should clear cookies and redirect to redirect_uri query param', async () => {
       jest.spyOn(authenticator._jwtVerifier, 'verify');
       authenticator._jwtVerifier.cacheJwks(jwksData);
@@ -501,7 +501,7 @@
       const response = await (authenticator as any)._clearCookies(request);
       expect(response).toEqual(expect.objectContaining({ status: '302' }));
       expect(response.headers['location']?.[0]?.value).toEqual('https://foobar.com');
     });
 
     it('should clear cookies and redirect to cf domain', async () => {
       jest.spyOn(authenticator._jwtVerifier, 'verify');
@@ -715,13 +715,38 @@
     const request = getCloudfrontRequest();
     request.Records[0].cf.request.querystring = 'code=54fe5f4e&state=/lol';
     return expect(authenticator.handle(request)).resolves.toEqual({ response: 'toto' })
       .then(() => {
         expect(authenticator._jwtVerifier.verify).toHaveBeenCalled();
         expect(authenticator._fetchTokensFromCode).toHaveBeenCalled();
         expect(authenticator._getRedirectResponse).toHaveBeenCalledWith(tokenData, 'd111111abcdef8.cloudfront.net', '/lol');
       });
   });
 
+  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);
@@ -761,13 +786,47 @@
             value: 'no-cache',
           }],
         },
       },
     )
       .then(() => {
         expect(authenticator._jwtVerifier.verify).toHaveBeenCalled();
       });
   });
 
+  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 = {
@@ -796,13 +855,28 @@
     expect(url.searchParams.get('state')).toBeDefined();
 
     // Cookies
     expect(response.headers['set-cookie']).toBeDefined();
     const cookies = response.headers['set-cookie'].map(h => h.value);
     expect(cookies.find(c => c.match(`.${NONCE_COOKIE_NAME_SUFFIX}=`))).toBeDefined();
     expect(cookies.find(c => c.match(`.${NONCE_HMAC_COOKIE_NAME_SUFFIX}=`))).toBeDefined();
     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 });

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,
+      );
     }
   }