AWS Cognito, Sign in with Apple and Native Apps

⛅ aws 🟦 typescript

As part of working on Storebar we have a requirement to authenticate users but have no interest in storing our user’s passwords ourselves. It’s a MacOS and iOS app so offering Sign in with Apple (SIWA) is a natural choice.

AWS Cognito is a good fit for us. It’s a managed service, low-cost and integrates nicely with the rest of the stack (specifically as an authoriser for API Gateway).

We quickly discovered that the out-of-the-box integration with SIWA was only supported on the web and didn’t allow native clients like MacOS apps authenticate in that way with the Cognito User Pool.

I’ll take you through how we managed to implement Sign in with Apple support on AWS Cognito with a Custom Authentication Challenge.

When the user clicks “Sign in with Apple” they are redirected to Apple’s authentication sheet page.

If that authentication is successful a JSON Web Token (JWT) is issued containing a stable ID to reference the user and other basic information like name and email address, if they were provided by the user.

If the JWT can be verified against Apple’s Keys then the payload is valid and the content can be trusted, allowing for the authentication of the user.

System Design Overview

System Design An overview of Storebar’s backend.

Our backend is very normal - an API that talks to the database. The specifics aren’t important now, but we exposed two auth endpoints: SIWA and OAuth.

These endpoints handle different scenarios: swapping a SIWA JWT for an access token and refreshing an access token respectively.

We decided not to integrate the AWS SDK directly into the apps but to instead do the server-to-server communication in the API itself.

Implementing Custom Authentication

The building blocks for extending Cognito are triggers. Triggers are Lambda Functions that are executed at specific points during authentication actions such as logging in or signing up.

Custom Authentication uses three triggers in particular:

The various custom authentication challenge triggers allow you to create any authentication system you need.

Define Challenge lets you look at the incoming request and decide what challenges you want to pose to the request.

In our case we just want to trigger our custom challenge. We also need to handle the case where the user has gone through the challenge flow already.

import type { DefineAuthChallengeTriggerHandler } from 'aws-lambda';

export const ChallengeDefine: DefineAuthChallengeTriggerHandler = async (event) => {
  event.response.issueTokens = false;
  event.response.failAuthentication = false;

  // Has a challenge already been completed?
  if (event.request.session.length > 0) {
    const challengeResult = event.request.session[0];

    if (challengeResult.challengeName === 'CUSTOM_CHALLENGE') {
      // Depending on the challenge result, issue tokens or fail.
      event.response.issueTokens = challengeResult.challengeResult;
      event.response.failAuthentication = !challengeResult.challengeResult;
    }
  }

  if (!event.response.failAuthentication && !event.response.issueTokens) {
    event.response.challengeName = 'CUSTOM_CHALLENGE';
  }

  return event;
};

Create Challenge is the point at which you can do what needs to be done server side to set your challenge up.

In our case we don’t need to do anything but this is where you could send an email, decide which security question to ask, etc. I don’t actually think this is necessary but added for completeness.

import type { CreateAuthChallengeTriggerHandler } from 'aws-lambda';

export const ChallengeCreate: CreateAuthChallengeTriggerHandler = async (event) => {
  return event;
};

Verify Challenge is where you check the user’s response to your challenge. This is where we need to do most of our work.

As you’ll recall, the response that we get to our challenge is a JWT from token. At this step we need to verify that the token is valid and return whether or not the “answer was correct”. This flag will be used as the challenge result back in the Define Challenge trigger.

import { z } from 'zod';
import { zu } from 'zod_utilz';
import { Logger } from '@aws-lambda-powertools/logger';
import type { VerifyAuthChallengeResponseTriggerHandler } from 'aws-lambda';

const IdentityTokenSchema = z.object({
  token: z.string(),
  identifier: z.string(),
  email: z.string().optional()
});

const logger = new Logger({ serviceName: 'challenge:verify' });

export const ChallengeVerify: VerifyAuthChallengeResponseTriggerHandler = async (event) => {
  event.response.answerCorrect = false;

  const appleIdentityToken = event.request.challengeAnswer;

  const payload = IdentityTokenSchema.parse(zu.stringToJSON().parse(appleIdentityToken));

  // Ensure that the payload matches the current user context. This is to
  // prevent any valid Apple JWT being used to log a user in.
  if (payload.identifier !== event.userName) {
    logger.warn('payload does not match cognito user name', {
      expected: event.userName,
      actual: payload.identifier
    });
    return event;
  }

  try {
    await verifyToken(payload);
    event.response.answerCorrect = true;
  } catch (e) {
    logger.error('invalid jwt', { error: e });
  }

  return event;
};

To dive into the actual verification itself, we use a few libraries for verifying the JWT and interacting with remote JSON Web Key Sets (JWKS).

import { JwksClient } from 'jwks-rsa';
import jwt from 'jsonwebtoken';

const appleKeysClient = new JwksClient({
  jwksUri: 'https://appleid.apple.com/auth/keys'
});

const getPublicKey = async (kid: string) => appleKeysClient.getSigningKey(kid).then((key) => key.getPublicKey());

export const verifyToken = async (token: IdentityToken, now?: DateTime): Promise<jwt.Jwt> => {
  const options = { complete: true, clockTimestamp: now?.toSeconds() } as const;
  const decoded = jwt.decode(token.token, options);

  if (decoded) {
    const keyId = decoded.header.kid;

    if (keyId) {
      const key = await getPublicKey(keyId);

      if (key) {
        return jwt.verify(token.token, key, {
          ...options,
          audience: ['your.bundle.identifier', 'your.other.bundle.identifier'],
          issuer: 'https://appleid.apple.com',
          subject: token.identifier,
          ignoreExpiration: false,
          ignoreNotBefore: false
        });
      } else {
        logger.error(`key with id ${keyId} not found`, { token });
        throw new Error('key not found');
      }
    } else {
      logger.error(`key id not defined in token`, { token });
      throw new Error('keyid not defined');
    }
  } else {
    logger.error(`unable to decode jwt`, { token });
    throw new Error('unable to decode jwt');
  }
};

Pre Sign Up Trigger The “Pre sign-up” trigger happens before the request hits the user pool.

One final trigger we need to implement is the Pre Sign-up Trigger. This isn’t a part of the custom authentication but it does allow us to automatically confirm the user during sign up.

import type { PreSignUpTriggerHandler } from 'aws-lambda';

export const PreSignUp: PreSignUpTriggerHandler = async (event) => {
  event.response.autoConfirmUser = true;
  return event;
};

Interacting with Cognito

Now that everything is set up on Cognito’s side we’re ready to interact with Cognito’s APIs and log our user in. The actual implementation isn’t that important here but the calls using the AWS SDK are.

It’s not the cleanest approach but it’s a working example! I assume that the user is logging in. If it turns out they weren’t then I register then and then log them in.

import { CLIENT_METADATA_KEY, IdentityToken, IdentityTokenSchema } from '@/lib/apple';
import { Logger } from '@aws-lambda-powertools/logger';
import {
  AuthFlowType,
  AuthenticationResultType,
  ChallengeNameType,
  CognitoIdentityProviderClient,
  InitiateAuthCommand,
  InitiateAuthCommandOutput,
  NotAuthorizedException,
  RespondToAuthChallengeCommand,
  RespondToAuthChallengeCommandOutput,
  SignUpCommand,
  UserNotFoundException
} from '@aws-sdk/client-cognito-identity-provider';
import { APIGatewayProxyHandlerV2 } from 'aws-lambda';
import { zu } from 'zod_utilz';
import crypto from 'crypto';
import { loadConfig } from '@/lib/config';

class EmailRequiredException extends Error {}

const client = new CognitoIdentityProviderClient({});

const LoginSIWA: APIGatewayProxyHandlerV2 = async (event) => {
  const config = await loadConfig();
  const payload = IdentityTokenSchema.parse(zu.stringToJSON().parse(event.body));

  try {
    try {
      // Assume that the user already has an account.
      const loginResponse = await client.send(createAuthCommand(config.auth.cognito.clientId, payload));

      // Send the payload as the challenge response.
      const challengeResponse = await client.send(
        createChallengeResponseCommand(config.auth.cognito.clientId, payload, loginResponse)
      );

      if (challengeResponse.AuthenticationResult) {
        return {
          statusCode: 200,
          body: JSON.stringify({
            data: formatAuthResult(challengeResponse.AuthenticationResult)
          })
        };
      } else {
        throw new Error('expected authentication result');
      }
    } catch (e) {
      // If the user is logging in for the first time then sign them up
      // and login again.
      if (e instanceof UserNotFoundException) {
        if (!payload.email) {
          throw new EmailRequiredException();
        }

        const signupResponse = await client.send(createSignUpCommand(config.auth.cognito.clientId, payload));

        const loginResponseAgain = await client.send(createAuthCommand(config.auth.cognito.clientId, payload));

        const challengeResponse = await client.send(
          createChallengeResponseCommand(config.auth.cognito.clientId, payload, loginResponseAgain)
        );

        if (challengeResponse.AuthenticationResult) {
          return {
            statusCode: 200,
            body: JSON.stringify({
              data: formatAuthResult(challengeResponse.AuthenticationResult)
            })
          };
        } else {
          throw new Error('expected authentication result');
        }
      }

      throw e;
    }
  } catch (e) {
    // handle your exceptions
  }
};

const createAuthCommand = (clientId: string, payload: IdentityToken) =>
  new InitiateAuthCommand({
    AuthFlow: AuthFlowType.CUSTOM_AUTH,
    AuthParameters: {
      USERNAME: payload.identifier
    },
    ClientId: clientId,
    ClientMetadata: { [CLIENT_METADATA_KEY]: JSON.stringify(payload) }
  });

const createSignUpCommand = (clientId: string, payload: IdentityToken) =>
  new SignUpCommand({
    ClientId: clientId,
    Username: payload.identifier,
    // Set a random password on the user.
    Password: crypto.randomBytes(64).toString('hex') + 'aB1!',
    ClientMetadata: { [CLIENT_METADATA_KEY]: JSON.stringify(payload) },
    UserAttributes: [
      {
        Name: 'email',
        Value: payload.email ?? ''
      }
    ]
  });

const createChallengeResponseCommand = (
  clientId: string,
  payload: IdentityToken,
  loginResponse: InitiateAuthCommandOutput
) =>
  new RespondToAuthChallengeCommand({
    ChallengeName: ChallengeNameType.CUSTOM_CHALLENGE,
    ClientId: clientId,
    ChallengeResponses: {
      ANSWER: JSON.stringify(payload),
      USERNAME: payload.identifier
    },
    Session: loginResponse.Session
  });

const formatAuthResult = (challengeResponse: AuthenticationResultType) => ({
  accessToken: challengeResponse.AccessToken,
  refreshToken: challengeResponse.RefreshToken,
  expiresIn: challengeResponse.ExpiresIn,
  tokenType: challengeResponse.TokenType,
  idToken: challengeResponse.IdToken
});

Bonus: Deploying with AWS CDK

import { UserPool } from 'aws-cdk-lib/aws-cognito';

// Omitted the creation of the AuthDefineChallenge, AuthCreateAuthChallenge
// AuthVerifyAuthChallengeResponse and AuthPreSignUp function definitions!

const userPool = new UserPool(this, 'UserPool', {
  userPoolName: 'myuserpool',
  lambdaTriggers: {
    defineAuthChallenge: AuthDefineChallenge,
    createAuthChallenge: AuthCreateAuthChallenge,
    verifyAuthChallengeResponse: AuthVerifyAuthChallengeResponse,
    preSignUp: AuthPreSignUp
  }
});

AWS Cognito, Sign in with Apple and Native Apps

How does Sign in With Apple Work?

System Design Overview

Implementing Custom Authentication

Interacting with Cognito

Bonus: Deploying with AWS CDK

Nick Jones