Skip to content

Commit

Permalink
Fix/expire all (#290)
Browse files Browse the repository at this point in the history
* Expire all cookies upon sign out

* Tweak
  • Loading branch information
ottokruse authored Sep 13, 2024
1 parent 6d7e45d commit 078da4e
Showing 1 changed file with 108 additions and 71 deletions.
179 changes: 108 additions & 71 deletions src/lambda-edge/shared/shared.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ function getDefaultCookieSettings(props: {
};
}
throw new Error(
`Cannot determine default cookiesettings for ${props.mode} with compatibility ${props.compatibility}`
`Cannot determine default cookie settings for ${props.mode} with compatibility ${props.compatibility}`
);
}

Expand Down Expand Up @@ -301,6 +301,7 @@ export function getAmplifyCookieNames(
idTokenKey: `${keyPrefix}.${tokenUserName}.idToken`,
accessTokenKey: `${keyPrefix}.${tokenUserName}.accessToken`,
refreshTokenKey: `${keyPrefix}.${tokenUserName}.refreshToken`,
hostedUiKey: "amplify-signin-with-hostedUI",
};
}

Expand Down Expand Up @@ -360,21 +361,21 @@ export const generateCookieHeaders = {
param: GenerateCookieHeadersParam & {
tokens: { id: string; access: string; refresh: string };
}
) => _generateCookieHeaders({ ...param }),
) => _generateCookieHeaders({ ...param, scenario: "SIGN_IN" }),
refresh: (
param: GenerateCookieHeadersParam & {
tokens: { id: string; access: string };
}
) => _generateCookieHeaders({ ...param }),
) => _generateCookieHeaders({ ...param, scenario: "REFRESH" }),
refreshFailed: (param: GenerateCookieHeadersParam) =>
_generateCookieHeaders({ ...param, expireCookies: "REFRESH_TOKEN" }),
_generateCookieHeaders({ ...param, scenario: "REFRESH_FAILED" }),
signOut: (param: GenerateCookieHeadersParam) =>
_generateCookieHeaders({ ...param, expireCookies: "ALL" }),
_generateCookieHeaders({ ...param, scenario: "SIGN_OUT" }),
};

function _generateCookieHeaders(
param: GenerateCookieHeadersParam & {
expireCookies?: "ALL" | "REFRESH_TOKEN";
scenario: "SIGN_IN" | "SIGN_OUT" | "REFRESH" | "REFRESH_FAILED";
}
) {
/**
Expand All @@ -389,69 +390,102 @@ function _generateCookieHeaders(

const decodedIdToken = decodeToken(param.tokens.id);
const tokenUserName = decodedIdToken["cognito:username"];
const userData = JSON.stringify({
UserAttributes: [
{
Name: "sub",
Value: decodedIdToken["sub"],
},
{
Name: "email",
Value: decodedIdToken["email"],
},
],
Username: tokenUserName,
});

const cookies: Cookies = {};
let cookieNames:
| ReturnType<typeof getAmplifyCookieNames>
| ReturnType<typeof getElasticsearchCookieNames>;
if (param.cookieCompatibility === "amplify") {
cookieNames = getAmplifyCookieNames(param.clientId, tokenUserName);
const userData = JSON.stringify({
UserAttributes: [
{
Name: "sub",
Value: decodedIdToken["sub"],
},
{
Name: "email",
Value: decodedIdToken["email"],
},
],
Username: tokenUserName,
});

// Construct object with the cookies
Object.assign(cookies, {
[cookieNames.lastUserKey]: `${tokenUserName}; ${param.cookieSettings.idToken}`,
[cookieNames.scopeKey]: `${param.oauthScopes.join(" ")}; ${
param.cookieSettings.accessToken
}`,
[cookieNames.userDataKey]: `${encodeURIComponent(userData)}; ${
param.cookieSettings.idToken
}`,
"amplify-signin-with-hostedUI": `true; ${param.cookieSettings.accessToken}`,
});
} else {
cookieNames = getElasticsearchCookieNames();
cookies[
cookieNames.cognitoEnabledKey
] = `True; ${param.cookieSettings.cognitoEnabled}`;
}

// Set JWTs in the cookies
cookies[
cookieNames.idTokenKey
] = `${param.tokens.id}; ${param.cookieSettings.idToken}`;
if (param.tokens.access) {
cookies[
const cookiesToSetOrExpire: Cookies = {};
const cookieNames =
param.cookieCompatibility === "amplify"
? getAmplifyCookieNames(param.clientId, tokenUserName)
: getElasticsearchCookieNames();

// Set or clear JWTs from the cookies
if (param.scenario === "SIGN_IN") {
// JWTs:
cookiesToSetOrExpire[
cookieNames.idTokenKey
] = `${param.tokens.id}; ${param.cookieSettings.idToken}`;
cookiesToSetOrExpire[
cookieNames.accessTokenKey
] = `${param.tokens.access}; ${param.cookieSettings.accessToken}`;
}
if (param.tokens.refresh) {
cookies[
cookiesToSetOrExpire[
cookieNames.refreshTokenKey
] = `${param.tokens.refresh}; ${param.cookieSettings.refreshToken}`;
}

if (param.expireCookies === "ALL") {
// Expire all cookies
Object.keys(cookies).forEach(
(key) => (cookies[key] = expireCookie(cookies[key]))
// Other cookies:
if ("lastUserKey" in cookieNames)
cookiesToSetOrExpire[
cookieNames.lastUserKey
] = `${tokenUserName}; ${param.cookieSettings.idToken}`;
if ("scopeKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.scopeKey] = `${param.oauthScopes.join(
" "
)}; ${param.cookieSettings.accessToken}`;
if ("userDataKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.userDataKey] = `${encodeURIComponent(
userData
)}; ${param.cookieSettings.idToken}`;
if ("hostedUiKey" in cookieNames)
cookiesToSetOrExpire[
cookieNames.hostedUiKey
] = `true; ${param.cookieSettings.accessToken}`;
if ("cognitoEnabledKey" in cookieNames)
cookiesToSetOrExpire[
cookieNames.cognitoEnabledKey
] = `True; ${param.cookieSettings.cognitoEnabled}`;
} else if (param.scenario === "REFRESH") {
cookiesToSetOrExpire[
cookieNames.idTokenKey
] = `${param.tokens.id}; ${param.cookieSettings.idToken}`;
cookiesToSetOrExpire[
cookieNames.accessTokenKey
] = `${param.tokens.access}; ${param.cookieSettings.accessToken}`;
} else if (param.scenario === "SIGN_OUT") {
// Expire JWTs
cookiesToSetOrExpire[cookieNames.idTokenKey] = addExpiry(
param.cookieSettings.idToken
);
} else if (param.expireCookies === "REFRESH_TOKEN") {
// Expire refresh token
cookies[cookieNames.refreshTokenKey] = expireCookie(
cookieNames.refreshTokenKey
cookiesToSetOrExpire[cookieNames.accessTokenKey] = addExpiry(
param.cookieSettings.accessToken
);
cookiesToSetOrExpire[cookieNames.refreshTokenKey] = addExpiry(
param.cookieSettings.refreshToken
);
// Expire other cookies
if ("lastUserKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.lastUserKey] = addExpiry(
param.cookieSettings.idToken
);
if ("scopeKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.scopeKey] = addExpiry(
param.cookieSettings.accessToken
);
if ("userDataKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.userDataKey] = addExpiry(
param.cookieSettings.idToken
);
if ("hostedUiKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.hostedUiKey] = addExpiry(
param.cookieSettings.accessToken
);
if ("cognitoEnabledKey" in cookieNames)
cookiesToSetOrExpire[cookieNames.cognitoEnabledKey] = addExpiry(
param.cookieSettings.cognitoEnabled
);
} else if (param.scenario === "REFRESH_FAILED") {
// Expire refresh token only
cookiesToSetOrExpire[cookieNames.refreshTokenKey] = addExpiry(
param.cookieSettings.refreshToken
);
}

Expand All @@ -461,31 +495,34 @@ function _generateCookieHeaders(
"spa-auth-edge-nonce-hmac",
"spa-auth-edge-pkce",
].forEach((key) => {
cookies[key] = expireCookie(`;${param.cookieSettings.nonce}`);
cookiesToSetOrExpire[key] = addExpiry(param.cookieSettings.nonce);
});

// Return cookie object in format of CloudFront headers
return Object.entries({
...param.additionalCookies,
...cookies,
...cookiesToSetOrExpire,
}).map(([k, v]) => ({ key: "set-cookie", value: `${k}=${v}` }));
}

function expireCookie(cookie: string = "") {
const cookieParts = cookie
/**
* Expire a cookie by setting its expiration time to the epoch start
* @param cookieSettings The cookie settings to add the expire setting to, for example: "Domain=example.com; Secure; HttpOnly"
* @returns Updated cookie settings that you can use as cookie value, i.e. with leading ; and expire instruction, for example: "; Domain=example.com; Secure; HttpOnly; Expires=Thu, 01 Jan 1970 00:00:00 GMT"
*/
function addExpiry(cookieSettings: string) {
const parts = cookieSettings
.split(";")
.map((part) => part.trim())
.filter((part) => !part.toLowerCase().startsWith("max-age"))
.filter((part) => !part.toLowerCase().startsWith("expires"));
const expires = `Expires=${new Date(0).toUTCString()}`;
const [, ...settings] = cookieParts; // first part is the cookie value, which we'll clear
return ["", ...settings, expires].join("; ");
return ["", ...parts, expires].join("; ");
}

function decodeToken(jwt: string) {
const tokenBody = jwt.split(".")[1];
const decodableTokenBody = tokenBody.replace(/-/g, "+").replace(/_/g, "/");
return JSON.parse(Buffer.from(decodableTokenBody, "base64").toString());
return JSON.parse(Buffer.from(tokenBody, "base64url").toString());
}

const AGENT = new Agent({ keepAlive: true });
Expand Down

0 comments on commit 078da4e

Please sign in to comment.