Fix: handle expiration of refresh token

src/lambda-edge/check-auth/index.ts

@@ -16,13 +16,15 @@ export const handler: CloudFrontRequestHandler = async (event) => {
   const requestedUri = `${request.uri}${
     request.querystring ? "?" + request.querystring : ""
   }`;
+  let refreshToken: string | undefined = "";
   try {
     const cookies = common.extractAndParseCookies(
       request.headers,
       CONFIG.clientId,
       CONFIG.cookieCompatibility
     );
     CONFIG.logger.debug("Extracted cookies:", cookies);
+    refreshToken = cookies.refreshToken;
 
     // If there's no ID token in your cookies, then you are not signed in yet
     if (!cookies.idToken) {
@@ -42,7 +44,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
     // If the JWT is expired we can try to refresh it
     // This is done by redirecting the user to the refresh path.
     // If the refresh works, the user will be redirected back here (this time with valid JWTs)
-    if (err instanceof common.JwtExpiredError) {
+    if (err instanceof common.JwtExpiredError && refreshToken) {
       CONFIG.logger.debug("Redirecting user to refresh path");
       return redirectToRefreshPath({ domainName, requestedUri });
     }
@@ -122,7 +124,7 @@ function redirectToCognitoHostedUI({
       ...CONFIG.cloudFrontHeaders,
     },
   };
-  CONFIG.logger.debug("Returning response:\n", response);
+  CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
   return response;
 }
 
@@ -150,7 +152,7 @@ function redirectToRefreshPath({
       ...CONFIG.cloudFrontHeaders,
     },
   };
-  CONFIG.logger.debug("Returning response:\n", response);
+  CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
   return response;
 }
 
@@ -182,7 +184,7 @@ function showContactAdminErrorPage({
       ],
     },
   };
-  CONFIG.logger.debug("Returning response:\n", response);
+  CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
   return response;
 }
 

src/lambda-edge/http-headers/index.ts

@@ -11,6 +11,6 @@ export const handler: CloudFrontResponseHandler = async (event) => {
   CONFIG.logger.debug("Event:", event);
   const response = event.Records[0].cf.response;
   Object.assign(response.headers, CONFIG.cloudFrontHeaders);
-  CONFIG.logger.debug("Returning response:\n", response);
+  CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
   return response;
 };

src/lambda-edge/parse-auth/index.ts

@@ -79,14 +79,6 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         );
       });
     CONFIG.logger.info("Successfully exchanged authorization code for tokens");
-    CONFIG.logger.debug("Response from Cognito token endpoint:\n", {
-      status,
-      headers,
-      idToken,
-      accessToken,
-      refreshToken,
-    });
-
     const response = {
       status: "307",
       statusDescription: "Temporary Redirect",
@@ -108,7 +100,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         ...CONFIG.cloudFrontHeaders,
       },
     };
-    CONFIG.logger.debug("Returning response:\n", response);
+    CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
     return response;
   } catch (err) {
     CONFIG.logger.error(err);
@@ -133,7 +125,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
           ...CONFIG.cloudFrontHeaders,
         },
       };
-      CONFIG.logger.debug("Returning response:\n", response);
+      CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
       return response;
     }
     let htmlParams: Parameters<typeof common.createErrorHtml>[0];
@@ -169,7 +161,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         ],
       },
     };
-    CONFIG.logger.debug("Returning response:\n", response);
+    CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
     return response;
   }
 };

src/lambda-edge/refresh-auth/index.ts

@@ -15,28 +15,25 @@ export const handler: CloudFrontRequestHandler = async (event) => {
   CONFIG.logger.debug("Event:", event);
   const request = event.Records[0].cf.request;
   const domainName = request.headers["host"][0].value;
+  let requestedUri: string | string[] | undefined = "/";
+  let idToken: string | undefined = undefined;
 
   try {
-    const { requestedUri, nonce: currentNonce } = parseQueryString(
-      request.querystring
-    );
-    const {
-      idToken,
-      refreshToken,
-      nonce: originalNonce,
-      nonceHmac,
-    } = common.extractAndParseCookies(
+    const querySting = parseQueryString(request.querystring);
+    requestedUri = querySting.requestedUri;
+    const cookies = common.extractAndParseCookies(
       request.headers,
       CONFIG.clientId,
       CONFIG.cookieCompatibility
     );
+    idToken = cookies.idToken;
 
     validateRefreshRequest(
-      currentNonce,
-      nonceHmac,
-      originalNonce,
-      idToken,
-      refreshToken
+      querySting.nonce,
+      cookies.nonceHmac,
+      cookies.nonce,
+      cookies.idToken,
+      cookies.refreshToken
     );
 
     const headers: { "Content-Type": string; Authorization?: string } = {
@@ -55,7 +52,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
     const body = stringifyQueryString({
       grant_type: "refresh_token",
       client_id: CONFIG.clientId,
-      refresh_token: refreshToken,
+      refresh_token: cookies.refreshToken,
     });
     const res = await common
       .httpPostToCognitoWithRetry(
@@ -67,6 +64,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
       .catch((err) => {
         throw new Error(`Failed to refresh tokens: ${err}`);
       });
+    CONFIG.logger.info("Successfully renewed tokens");
     newIdToken = res.data.id_token as string;
     newAccessToken = res.data.access_token as string;
     const response = {
@@ -88,9 +86,46 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         ...CONFIG.cloudFrontHeaders,
       },
     };
-    CONFIG.logger.debug("Returning response:\n", response);
+    CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
     return response;
   } catch (err) {
+    if (
+      err instanceof Error &&
+      err.message.includes("invalid_grant") &&
+      idToken
+    ) {
+      // The refresh token has likely expired.
+      // We'll clear the refresh token cookie, so that CheckAuth won't redirect any more requests here in vain.
+      // Also, we'll redirect the user to where he/she came from.
+      // (From there CheckAuth will redirect the user to the Cognito hosted UI to sign in)
+      CONFIG.logger.info(
+        "Expiring refresh token cookie, as the refresh token has expired"
+      );
+      const response = {
+        status: "307",
+        statusDescription: "Temporary Redirect",
+        headers: {
+          location: [
+            {
+              key: "location",
+              value: `https://${domainName}${common.ensureValidRedirectPath(
+                requestedUri
+              )}`,
+            },
+          ],
+          "set-cookie": common.generateCookieHeaders.refreshFailed({
+            tokens: {
+              id: idToken,
+            },
+            ...CONFIG,
+          }),
+          ...CONFIG.cloudFrontHeaders,
+        },
+      };
+      CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
+      return response;
+    }
+    CONFIG.logger.error(err);
     const response = {
       body: common.createErrorHtml({
         title: "Refresh issue",
@@ -111,7 +146,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         ],
       },
     };
-    CONFIG.logger.debug("Returning response:\n", response);
+    CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
     return response;
   }
 };

src/lambda-edge/shared/shared.ts

@@ -360,28 +360,31 @@ export const generateCookieHeaders = {
     param: GenerateCookieHeadersParam & {
       tokens: { id: string; access: string; refresh: string };
     }
-  ) => _generateCookieHeaders({ ...param, event: "signIn" }),
+  ) => _generateCookieHeaders({ ...param }),
   refresh: (
     param: GenerateCookieHeadersParam & {
       tokens: { id: string; access: string };
     }
-  ) => _generateCookieHeaders({ ...param, event: "refresh" }),
+  ) => _generateCookieHeaders({ ...param }),
+  refreshFailed: (param: GenerateCookieHeadersParam) =>
+    _generateCookieHeaders({ ...param, expireCookies: "REFRESH_TOKEN" }),
   signOut: (param: GenerateCookieHeadersParam) =>
-    _generateCookieHeaders({ ...param, event: "signOut" }),
+    _generateCookieHeaders({ ...param, expireCookies: "ALL" }),
 };
 
 function _generateCookieHeaders(
   param: GenerateCookieHeadersParam & {
-    event: "signIn" | "signOut" | "refresh";
+    expireCookies?: "ALL" | "REFRESH_TOKEN";
   }
 ) {
   /**
-   * Generate cookie headers for the following scenario's:
-   *  - signIn: called from Parse Auth lambda, when receiving fresh JWTs from Cognito
-   *  - sign out: called from Sign Out Lambda, when the user visits the sign out URL
-   *  - refresh: called from Refresh Auth lambda, when receiving fresh ID and Access JWTs from Cognito
+   * Generate cookie headers to set, or clear, cookies with JWTs.
    *
-   *   Note that there are other places besides this helper function where cookies can be set (search codebase for "set-cookie")
+   * This is centralized in this function because there is logic to determine
+   * the right cookie names, that we do not want to repeat everywhere.
+   *
+   * Note that there are other places besides this helper function where
+   * cookies can be set (search codebase for "set-cookie").
    */
 
   const decodedIdToken = decodeToken(param.tokens.id);
@@ -440,17 +443,19 @@ function _generateCookieHeaders(
     ] = `${param.tokens.refresh}; ${param.cookieSettings.refreshToken}`;
   }
 
-  if (param.event === "signOut") {
+  if (param.expireCookies === "ALL") {
     // Expire all cookies
     Object.keys(cookies).forEach(
       (key) => (cookies[key] = expireCookie(cookies[key]))
     );
+  } else if (param.expireCookies === "REFRESH_TOKEN") {
+    // Expire refresh token
+    cookies[cookieNames.refreshTokenKey] = expireCookie(
+      cookieNames.refreshTokenKey
+    );
   }
 
-  // Always expire nonce, nonceHmac and pkce - this is valid in all scenario's:
-  // * event === 'newTokens' --> you just signed in and used your nonce and pkce successfully, don't need them no more
-  // * event === 'refreshFailed' --> you are signed in already, why do you still have a nonce?
-  // * event === 'signOut' --> clear ALL cookies anyway
+  // Always expire nonce, nonceHmac and pkce
   [
     "spa-auth-edge-nonce",
     "spa-auth-edge-nonce-hmac",
@@ -485,6 +490,8 @@ function decodeToken(jwt: string) {
 
 const AGENT = new Agent({ keepAlive: true });
 
+class NonRetryableFetchError extends Error {}
+
 export async function httpPostToCognitoWithRetry(
   url: string,
   data: Buffer,
@@ -500,22 +507,41 @@ export async function httpPostToCognitoWithRetry(
         ...options,
         method: "POST",
       }).then((res) => {
-        if (res.status !== 200) {
-          throw new Error(`Status is ${res.status}, expected 200`);
-        }
+        const responseData = res.data.toString();
+        logger.debug(
+          `Response from Cognito:`,
+          JSON.stringify({
+            status: res.status,
+            headers: res.headers,
+            data: responseData,
+          })
+        );
         if (!res.headers["content-type"]?.startsWith("application/json")) {
           throw new Error(
             `Content-Type is ${res.headers["content-type"]}, expected application/json`
           );
         }
+        const parsedResponseData = JSON.parse(responseData);
+        if (res.status !== 200) {
+          const errorMessage =
+            parsedResponseData.error || `Status is ${res.status}, expected 200`;
+          if (res.status && res.status >= 400 && res.status < 500) {
+            // No use in retrying client errors
+            throw new NonRetryableFetchError(errorMessage);
+          } else {
+            throw new Error(errorMessage);
+          }
+        }
         return {
           ...res,
-          data: JSON.parse(res.data.toString()),
+          data: parsedResponseData,
         };
       });
     } catch (err) {
-      logger.debug(`HTTP POST to ${url} failed (attempt ${attempts}):`);
-      logger.debug(err);
+      logger.debug(`HTTP POST to ${url} failed (attempt ${attempts}): ${err}`);
+      if (err instanceof NonRetryableFetchError) {
+        throw err;
+      }
       if (attempts >= 5) {
         // Try 5 times at most
         logger.error(

src/lambda-edge/sign-out/index.ts

@@ -20,13 +20,13 @@ export const handler: CloudFrontRequestHandler = async (event) => {
   CONFIG.logger.debug("Event:", event);
   const request = event.Records[0].cf.request;
   const domainName = request.headers["host"][0].value;
-  const { idToken, accessToken, refreshToken } = extractAndParseCookies(
+  const cookies = extractAndParseCookies(
     request.headers,
     CONFIG.clientId,
     CONFIG.cookieCompatibility
   );
 
-  if (!idToken) {
+  if (!cookies.idToken) {
     const response = {
       body: createErrorHtml({
         title: "Signed out",
@@ -45,7 +45,7 @@ export const handler: CloudFrontRequestHandler = async (event) => {
         ],
       },
     };
-    CONFIG.logger.debug("Returning response:\n", response);
+    CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
     return response;
   }
 
@@ -68,13 +68,13 @@ export const handler: CloudFrontRequestHandler = async (event) => {
       ],
       "set-cookie": generateCookieHeaders.signOut({
         tokens: {
-          id: idToken,
+          id: cookies.idToken,
         },
         ...CONFIG,
       }),
       ...CONFIG.cloudFrontHeaders,
     },
   };
-  CONFIG.logger.debug("Returning response:\n", response);
+  CONFIG.logger.debug("Returning response:\n", JSON.stringify(response));
   return response;
 };