Simple federated sign-on with Amazon Cognito Part 2 - The code

Now that we've got the general setup out of the way in part 1, it's time to dig into how the cognito.js code actually works. Keep in mind it's dependent on js-sha256 for the SHA256 implementation, which is included for you if you use the example index.html file.




First, we define the constants used by the code:

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

//Constants we need const APP_BASE_URL = '<APP_BASE_URL>'; const AWS_COGNITO_POOL_CLIENT_ID = '<COGNITO_POOL_CLIENT_ID>'; const AWS_COGNITO_POOL_URI = '<COGNITO_POOL_URI>'; //The following must be registered with the Cognito app client const AWS_COGNITO_LOGIN_CALLBACK_URI = APP_BASE_URL + '/?action=login'; const AWS_COGNITO_LOGOUT_CALLBACK_URI = APP_BASE_URL + '/?action=logout'; //API endpoints const AWS_COGNITO_AUTHORIZE_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/authorize'; const AWS_COGNITO_CODE_EXCHANGE_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/token'; const AWS_COGNITO_USER_INFO_ENDPOINT = AWS_COGNITO_POOL_URI + '/oauth2/userInfo'; const AWS_COGNITO_LOGOUT_ENDPOINT = AWS_COGNITO_POOL_URI + '/logout?client_id=' + AWS_COGNITO_POOL_CLIENT_ID + '&logout_uri=' + AWS_COGNITO_LOGOUT_CALLBACK_URI; //Somewhere to store auth tokens, we don't use cookies or local storage because we have no way of knowing if the user has selected 'remember me' or not, so we only 'remember' for the length of a session const sessionStorage = window.sessionStorage;

'APP_BASE_URL' is the base URL for the page, likely 'http://localhost:8000' if you're serving the example using python. 'AWS_COGNITO_POOL_CLIENT_ID' is the client id of your Cognito user pool, 'AWS_COGNITO_URI' is the URI of your Cognito user pool.

The rest of the necessary constants are derived, and should only be changed if your program flow is different than the one described in part 1. 'AWS_COGNITO_LOGIN_CALLBACK_URI' is the URI we will return to after an authorization request (after a request to the AUTHORIZATION endpoint), we return here whether the request succeeded or failed. 'AWS_COGNITO_LOGOUT_CALLBACK_URI' is the URI returned to after a logout request (a request to the LOGOUT endpoint). Our example application is effectively a single page application, so the same URL handles both with the only difference being query parameters. If your application has distinct login and logout pages these URIs will have to be modified to match. Remember that these URIs will have to be registered with the Cognito app client.

Next up with define the OAuth2 endpoints as implemented by Amazon Cognito. The AUTHORIZATION endpoint is used over the LOGIN endpoint because the AUTHORIZATION endpoint explicitly supports PKCE. While my experiments show PKCE works with the LOGIN endpoint, it is not explicitly shown to be supported in the Cognito documentation.

Finally, we grab a reference to the 'sessionStorage'. We could store any relevant session tokens as cookies, but we want to keep our tokens as secure as possible and cookies are passed with requests. We use window.sessionStorage over window.localStorage because the refresh cycle can only be gone through once, so it doesn't make much sense to maintain tokens indefinitely. Instead, we let our tokens be deleted with the page session.

Next up, we decide what to do when the page is loaded (Should this be encapsulated in a window.onload event? Almost certainly, does it need to be? Apparently not):

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32

//Parse the URL params const urlParams = new URL(window.location.href).searchParams; const action = urlParams.get('action'); if (action === 'login') { if (sessionStorage.getItem('cognito_login_state') && sessionStorage.getItem('cognito_login_state') === urlParams.get('state') && sessionStorage.getItem('cognito_pkce_code_verifier')) { //Clear the state variable, it shouldn't be reused sessionStorage.removeItem('cognito_login_state'); //Exchange the auth token for an access token, get the current user const authCode = urlParams.get('code'); exchangeAuthToken(authCode, sessionStorage.getItem('cognito_pkce_code_verifier'), tokenSuccessCallback, errorCallback); } } else if (action === 'logout') { clearAuthData(); //Redirect to base URL window.location.href = APP_BASE_URL; } else if (sessionStorage.getItem('cognito_expiration') && sessionStorage.getItem('cognito_access_token')) { if (Math.floor(Date.now() / 1000) < sessionStorage.getItem('cognito_expiration')) { //We have a current token getCurrentUser(sessionStorage.getItem('cognito_access_token'), getUserSuccessCallback, errorCallback); } else { //Token is expired, clean up now defunct auth data clearAuthData(); } }

First, we parse the query parameters into 'urlParams' so we can use them. Technically this method isn't supported by Internet Explorer, you've been warned. From there, we parse the 'action' query parameter into the 'action' variable. The only actions we should ever see are 'login' and 'logout', defined in the login and logout callback URIs.

If the action is 'login' and we have the necessary parts to continue with authorization (a 'state' variable and a PKCE code verifier) start the process to exchange the given authorization code for an access token by calling 'exchangeAuthToken'. If we don't have a PKCE code verifier set, we don't continue as we'd have no way to prove to Cognito that we are the progenitor of the original authorization request.  If the 'state' variable isn't set, or the set value doesn't match the value of the 'state' query parameter, we don't continue. This helps mitigate the risk of a CSRF attack by helping ensure the party that initiated the callback is the same party that we made the authorization request to.

If the action is 'logout', we clear any stored authorization values, and redirect to the base url.

If we have a stored access token that isn't expired, we retrieve the current user. If we have an access token and it's expired, we clear the now useless stored authorization values.

Now let's look at the authorization flow in more detail, starting with the click handler that's called when the user clicks "Login":

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

function loginClickFunction() { //Generate a 'state' variable we can use to make sure we're returning from an authentication flow we started, this adds no security, it's purely for flow validation const stateParameter = generateSecret(); //Store the state so we can verify it later sessionStorage.setItem('cognito_login_state', stateParameter); //Generate a PKCE code verifier, again, since this is a public client, it provides no security, use a hash because it's easy, and satisifes length requirements //From RFC 7636, the code verifier is base64url-encoded ASCII const codeVerifier = base64UrlEncode(generateSecret()); sessionStorage.setItem('cognito_pkce_code_verifier', codeVerifier); //The challenge, by definition, is base64 encoded SHA256 hash of the code verifier const codeChallenge = hexToBase64UrlEncode(sha256(codeVerifier)); //Build the request query const loginRequestPairs = []; loginRequestPairs.push(encodeURIComponent('client_id') + '=' + encodeURIComponent(AWS_COGNITO_POOL_CLIENT_ID)); loginRequestPairs.push(encodeURIComponent('redirect_uri') + '=' + encodeURIComponent(AWS_COGNITO_LOGIN_CALLBACK_URI)); loginRequestPairs.push(encodeURIComponent('response_type') + '=' + encodeURIComponent('code')); loginRequestPairs.push(encodeURIComponent('code_challenge') + '=' + encodeURIComponent(codeChallenge)); loginRequestPairs.push(encodeURIComponent('code_challenge_method') + '=' + encodeURIComponent('S256')); loginRequestPairs.push(encodeURIComponent('state') + '=' + encodeURIComponent(stateParameter)); const loginRequestQuery = loginRequestPairs.join('&').replace(/%20/g, '+'); //Redirect to login window.location.href = AWS_COGNITO_AUTHORIZE_ENDPOINT + '?' + loginRequestQuery; return false; }

First, we generate a secret and store in in the 'stateParameter' variable. 'generateSecret()' returns a SHA256 hash as a hexadecimal string. We then store this state as a session variable so we can verify it during a login action as we saw earlier. We then generate another secret stored in the 'codeVerifier' variable and also in the session. This is used as part of the PKCE proof key exchange, RFC 7636 specifies the "code verifier" must be base64url encoded ASCII (see Appendix A of the RFC for details of how that differs from URL encoding a base64 encoded ASCII string). To initiate a PKCE secured request, we also need a "code challenge" which by definition is a base64url encoded SHA256 hash of the previously determined 'codeVerifier', this is stored in the 'codeChallenge' variable, but is not stored as a session variable as there is no use for it after this request. Finally, we build the request in the form of a URL with the necessary components as query parameters, and then redirect the user there. Most notably, we see the 'redirect_uri' parameter which will return the user to our login callback URI defined earlier.

As mentioned before, upon returning from the AUTHORIZE endpoint, the action query parameter will be set to 'login' and assuming everything is in order, we call 'exchangeAuthToken':

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33

function exchangeAuthToken(authorizationCode, codeVerifier, successCallback, errorCallback) { //Exchanges an authorization code for a set of user tokens //https://docs.aws.amazon.com/cognito/latest/developerguide/token-endpoint.html const tokenRequest = new XMLHttpRequest(); //Build the form data const tokenRequestPairs = []; tokenRequestPairs.push(encodeURIComponent('grant_type') + '=' + encodeURIComponent('authorization_code')); tokenRequestPairs.push(encodeURIComponent('client_id') + '=' + encodeURIComponent(AWS_COGNITO_POOL_CLIENT_ID)); tokenRequestPairs.push(encodeURIComponent('redirect_uri') + '=' + encodeURIComponent(AWS_COGNITO_LOGIN_CALLBACK_URI)); tokenRequestPairs.push(encodeURIComponent('code_verifier') + '=' + encodeURIComponent(codeVerifier)); tokenRequestPairs.push(encodeURIComponent('code') + '=' + encodeURIComponent(authorizationCode)); const tokenRequestData = tokenRequestPairs.join('&').replace(/%20/g, '+'); //Set up event handler tokenRequest.onreadystatechange = function() { if (tokenRequest.readyState == 4) { if (tokenRequest.status == 200) { successCallback(JSON.parse(tokenRequest.responseText)); } else { errorCallback(JSON.parse(tokenRequest.responseText)); } } }; //Submit the form tokenRequest.open('POST', AWS_COGNITO_CODE_EXCHANGE_ENDPOINT, true); tokenRequest.setRequestHeader('Content-Type', 'application/x-www-form-urlencoded'); tokenRequest.send(tokenRequestData); }

'exchangeAuthToken' gets called with the value of 'authorizationCode' set to the value of the 'code' query parameter returned to the login URI, and 'codeVerifier' set to the value stored in the session by the login click handler. Nothing particularly fancy happens here, we build the request exactly as defined in the Cognito documentation for the TOKEN endpoint. Note that if the PKCE code verifier is incorrect somehow, a 'request_error' will be returned. This is in contradiction to the RFC which calls out an 'invalid_grant' MUST be returned.

Let's take a look at the success callback for 'exchangeAuthToken':

1 2 3 4 5 6 7 8 9 10 11 12 13

function tokenSuccessCallback(tokenData) { //Save the data sessionStorage.setItem('cognito_access_token', tokenData['access_token']); sessionStorage.setItem('cognito_id_token', tokenData['id_token']); sessionStorage.setItem('cognito_refresh_token', tokenData['refresh_token']); sessionStorage.setItem('cognito_expiration', Math.floor(Date.now() / 1000) + tokenData['expires_in']); //Don't reuse PKCE token sessionStorage.removeItem('cognito_pkce_code_verifier'); //Redirect to base URL, the user info request will be made there window.location.href = APP_BASE_URL; }

Pretty simple, we store the various authorization values in the session and clean up the now useless PKCE code verifier. We then redirect the user to the base URL, which, now that our session is populated with all the necessary parts, should call 'getCurrentUser':

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

function getCurrentUser(accessToken, successCallback, errorCallback) { //Gets the user info //https://docs.aws.amazon.com/cognito/latest/developerguide/userinfo-endpoint.html const userInfoRequest = new XMLHttpRequest(); //Setup handler userInfoRequest.onreadystatechange = function() { if (userInfoRequest.readyState == 4) { if (userInfoRequest.status == 200) { successCallback(JSON.parse(userInfoRequest.responseText)); } else { errorCallback(JSON.parse(tokenRequest.responseText)); } } }; userInfoRequest.open('GET', AWS_COGNITO_USER_INFO_ENDPOINT, true); userInfoRequest.setRequestHeader('Authorization', 'Bearer ' + accessToken); userInfoRequest.send(null); }

Another pretty simple one, exactly as described in the documentation for the USERINFO endpoint. Which brings us to the final success callback:

1 2 3 4 5 6 7 8 9

function getUserSuccessCallback(userData) { //Hide the login link document.getElementById('hrefLogin').style.display = 'none'; //Show the logout link document.getElementById('hrefLogout').style.display = 'inline'; console.log(userData); }

Hide the login link, show the logout link, and spit out the returned user data to the web console. This is also called on a refresh of the base url.

To logout, there's the 'logoutClickFunction':

1 2 3 4 5

function logoutClickFunction() { //Redirect to logout window.location.href = AWS_COGNITO_LOGOUT_ENDPOINT; return false; }

This redirects the user to the LOGOUT endpoint, which will return them to the base url with the 'logout' action as an url parameter, which in turn will clear the session tokens, and redirect to the base url.

That's essentially it, just some helper functions left. There's a generic error callback to print errors to the web console:

1 2 3

function errorCallback(errorData) { console.log(errorData); }

There's the 'base64UrlEncode' function which takes a string of ASCII characters and base64url encodes it as per Appendix A of RFC 7636:

1 2 3 4 5 6 7 8 9

function base64UrlEncode(toEncode) { var encodedString = btoa(toEncode); encodedString = encodedString.replace(/=+$/, ''); encodedString = encodedString.replace(/\+/g, '-'); encodedString = encodedString.replace(/\//g, '_'); return encodedString; }

There's also the 'hexToBase64UrlEncode' which takes a string of bytes as a hex string, converts it to raw bytes interpreted as an ASCII string, which is then returned as a base64url encoded string:

1 2 3 4 5 6 7 8 9 10 11

function hexToBase64UrlEncode(hexString) { //https://stackoverflow.com/a/30614192 hexString = hexString.replace(/\r|\n/g, ''); hexString = hexString.replace(/([\da-fA-F]{2}) ?/g, '0x$1 '); hexString = hexString.replace(/ +$/, ''); const hexArray = hexString.split(' '); const byteString = String.fromCharCode.apply(null, hexArray); return base64UrlEncode(byteString); }

And finally, 'generateSecret' which uses window.crypto to to generate our "secrets":

1 2 3 4

function generateSecret() { //Generates a cryptographically secure "secret" with 256 bits of entropy, returned as a hex string return sha256(window.crypto.getRandomValues(new Uint32Array(8))); }

8 Uint32 * 32 bits / Uint32 gives us 256 bits of entropy. However, our ability to keep these values 'secret' is questionable. Anyway, we generate them, and return them as hex strings.

Hopefully this gives you some insight into how this all works. It's not the simplest implementation, but it's use of of the authorization code grant type, state value, and PKCE should be applicable to a wide range of use cases.