Customizing Access Tokens

This was last updated with the following packages:

ampx info

System:
  OS: macOS 14.7
  CPU: (10) arm64 Apple M1 Pro
  Memory: 116.13 MB / 32.00 GB
  Shell: /opt/homebrew/bin/fish
Binaries:
  Node: 22.8.0 - ~/.local/state/fnm_multishells/2908_1728338201728/bin/node
  Yarn: undefined - undefined
  npm: 10.8.2 - ~/.local/state/fnm_multishells/2908_1728338201728/bin/npm
  pnpm: 9.10.0 - ~/.local/state/fnm_multishells/2908_1728338201728/bin/pnpm
NPM Packages:
  @aws-amplify/auth-construct: Not Found
  @aws-amplify/backend: 1.3.0
  @aws-amplify/backend-auth: Not Found
  @aws-amplify/backend-cli: 1.2.8
  @aws-amplify/backend-data: Not Found
  @aws-amplify/backend-deployer: Not Found
  @aws-amplify/backend-function: Not Found
  @aws-amplify/backend-output-schemas: Not Found
  @aws-amplify/backend-output-storage: Not Found
  @aws-amplify/backend-secret: Not Found
  @aws-amplify/backend-storage: Not Found
  @aws-amplify/cli-core: Not Found
  @aws-amplify/client-config: Not Found
  @aws-amplify/deployed-backend-client: Not Found
  @aws-amplify/form-generator: Not Found
  @aws-amplify/model-generator: Not Found
  @aws-amplify/platform-core: Not Found
  @aws-amplify/plugin-types: Not Found
  @aws-amplify/sandbox: Not Found
  @aws-amplify/schema-generator: Not Found
  aws-amplify: 6.6.2
  aws-cdk: 2.160.0
  aws-cdk-lib: 2.160.0
  typescript: 5.6.2
AWS environment variables:
  AWS_PROFILE = josef
  AWS_REGION = us-east-1
  AWS_STS_REGIONAL_ENDPOINTS = regional
  AWS_NODEJS_CONNECTION_REUSE_ENABLED = 1
  AWS_SDK_LOAD_CONFIG = 1
No CDK environment variables

Files

• Directoryamplify

  • Directoryauth

    • Directorypre-token-generation

      • handler.ts
      • resource.ts
    • resource.ts
  • Directorydata

    • resource.ts
  • create-auth-resource-server.ts
  • backend.ts
  • constants.ts
• package.json

amplify/auth/resource.ts

import { defineAuth } from "@aws-amplify/backend"
import { GROUPS } from "../constants"
import { preTokenGeneration } from "./pre-token-generation/resource"

/**
 * Define and configure your auth resource
 * @see https://docs.amplify.aws/gen2/build-a-backend/auth
 */
export const auth = defineAuth({
  loginWith: {
    email: true,
  },
  groups: Object.values(GROUPS),
  triggers: {
    // do not put it here
    // preTokenGeneration,
  },
})

amplify/constants.ts

export const GROUPS = {
  ADMINS: "ADMINS",
  EVERYONE: "EVERYONE",
} as const

amplify/backend.ts

import type { UserPoolClient } from "aws-cdk-lib/aws-cognito"
import {
  UserPool,
  UserPoolOperation,
  LambdaVersion,
} from "aws-cdk-lib/aws-cognito"
import { defineBackend } from "@aws-amplify/backend"
import { auth } from "./auth/resource"
import { preTokenGeneration } from "./auth/pre-token-generation/resource"
import { data } from "./data/resource"
import { createAuthResourceServer } from "./overrides/create-auth-resource-server"

/**
 * @see https://docs.amplify.aws/react/build-a-backend/ to add storage, functions, and more
 */
const backend = defineBackend({
  auth,
  data,
  preTokenGeneration,
})

// enable ASF
backend.auth.resources.cfnResources.cfnUserPool.userPoolAddOns = {
  advancedSecurityMode: "AUDIT",
}

// add the trigger with Lambda event version 2
// this will _always_ be true with defineAuth
if (backend.auth.resources.userPool instanceof UserPool) {
  backend.auth.resources.userPool.addTrigger(
    UserPoolOperation.PRE_TOKEN_GENERATION_CONFIG,
    backend.preTokenGeneration.resources.lambda,
    LambdaVersion.V2_0
  )
}

createAuthResourceServer(backend.auth.stack, {
  userPool: backend.auth.resources.userPool as UserPool,
  userPoolClient: backend.auth.resources.userPoolClient as UserPoolClient,
})

amplify/create-auth-resource-server.ts

import type { UserPool, UserPoolClient } from "aws-cdk-lib/aws-cognito"
import type { Construct } from "constructs"
import { CfnUserPoolClient, ResourceServerScope } from "aws-cdk-lib/aws-cognito"

type CreateAuthResourceServerProps = {
  userPool: UserPool
  userPoolClient: UserPoolClient
}

export function createAuthResourceServer(
  scope: Construct,
  props: CreateAuthResourceServerProps
) {
  const { userPool, userPoolClient } = props

  const adminFullAccessScope = new ResourceServerScope({
    scopeName: "*",
    scopeDescription: "allow all admin operations",
  })

  const adminResourceServer = userPool.addResourceServer(
    "AdminResourceServer",
    {
      identifier: "admin",
      scopes: [adminFullAccessScope],
    }
  )

  const everyoneFullAccessScope = new ResourceServerScope({
    scopeName: "*",
    scopeDescription: "allow all everyone operations",
  })

  const everyoneResourceServer = userPool.addResourceServer(
    "EveryoneResourceServer",
    {
      identifier: "everyone",
      scopes: [everyoneFullAccessScope],
    }
  )

  const cfnUserPoolClient = userPoolClient.node.defaultChild

  if (!(cfnUserPoolClient instanceof CfnUserPoolClient)) {
    throw new Error("invalid user pool client child")
  }

  cfnUserPoolClient.allowedOAuthScopes = [
    ...(cfnUserPoolClient.allowedOAuthScopes || []),
    `${adminResourceServer.userPoolResourceServerId}/${adminFullAccessScope.scopeName}`,
    `${everyoneResourceServer.userPoolResourceServerId}/${everyoneFullAccessScope.scopeName}`,
  ]
}

amplify/auth/pre-token-generation/handler.ts

import type { PreTokenGenerationTriggerHandler } from "aws-lambda"
import { GROUPS, SCOPES } from "../../constants"

/**
 * Add scopes to tokens based on group association
 * https://docs.aws.amazon.com/cognito/latest/developerguide/user-pool-lambda-pre-token-generation.html#aws-lambda-triggers-pre-token-generation-example-version-2-overview
 */
export const handler: PreTokenGenerationTriggerHandler = async (event) => {
  console.log("event", JSON.stringify(event, null, 2))
  const groups = event.request.groupConfiguration.groupsToOverride
  const scopes = []

  if (groups?.includes(GROUPS.EVERYONE)) {
    scopes.push("everyone/*")
  }

  if (groups?.includes(GROUPS.ADMINS)) {
    scopes.push("admin/*")
  }

  // @ts-expect-error types are bad
  event.response.claimsAndScopeOverrideDetails = {
    accessTokenGeneration: {
      scopesToAdd: scopes,
    },
  }

  return event
}

Sample Frontend

src/App.tsx

import type { AuthSession } from "aws-amplify/auth"
import type { Schema } from "../amplify/data/resource"
import { useState, useEffect } from "react"
import { Authenticator, useAuthenticator } from "@aws-amplify/ui-react"
import { Amplify } from "aws-amplify"
import { fetchAuthSession } from "aws-amplify/auth"
import outputs from "../amplify_outputs.json"
import "@aws-amplify/ui-react/styles.css"
import { generateClient } from "aws-amplify/api"

Amplify.configure(outputs)

const client = generateClient<Schema>()

function Main() {
  const [result, setResult] = useState<Record<string, unknown>>()
  const [selected, setSelected] = useState<"id-token" | "access-token">(
    "access-token"
  )
  const [session, setSession] = useState<AuthSession>()
  const { user, signOut } = useAuthenticator()

  useEffect(() => {
    fetchAuthSession().then(setSession)
  }, [])

  async function query(operation: keyof typeof client.queries) {
    const accessToken = session?.tokens?.accessToken.toString()
    // biome-ignore lint/style/noNonNullAssertion: types are funky, this will always be there
    const idToken = session?.tokens?.idToken!.toString()

    return client.queries[operation]({
      headers: {
        Authorization: `Bearer ${
          selected === "access-token" ? accessToken : idToken
        }`,
      },
    })
  }

  return (
    <main className="container mx-auto">
      <h1>hello {user.userId}</h1>
      <button type="button" onClick={() => signOut()}>
        sign out
      </button>
      <section>
        {session ? (
          <>
            <div>
              <p>current selection: {selected}</p>
            </div>
            <fieldset>
              <legend className="sr-only">Select token type</legend>
              <div className="grid grid-cols-2 gap-4">
                <div className="flex-1">
                  <input
                    type="radio"
                    id="id-token"
                    name="id-token"
                    value="id-token"
                    checked={selected === "id-token"}
                    onChange={() => setSelected("id-token")}
                    className="peer sr-only"
                  />
                  <label
                    htmlFor="id-token"
                    className="flex flex-col items-center justify-center p-4 h-[34rem] bg-white border-2 border-gray-200 rounded-lg cursor-pointer peer-checked:border-blue-600 peer-checked:text-blue-600 hover:bg-gray-50 transition-all"
                  >
                    <span className="text-lg font-semibold mb-2">ID Token</span>
                    <pre className="w-full h-full overflow-auto text-sm bg-gray-100 p-2 rounded">
                      {JSON.stringify(
                        session.tokens?.idToken?.payload,
                        null,
                        2
                      )}
                    </pre>
                  </label>
                </div>
                <div className="flex-1">
                  <input
                    type="radio"
                    id="access-token"
                    name="access-token"
                    value="access-token"
                    checked={selected === "access-token"}
                    onChange={() => setSelected("access-token")}
                    className="peer sr-only"
                  />
                  <label
                    htmlFor="access-token"
                    className="flex flex-col items-center justify-center p-4 h-[34rem] bg-white border-2 border-gray-200 rounded-lg cursor-pointer peer-checked:border-blue-600 peer-checked:text-blue-600 hover:bg-gray-50 transition-all"
                  >
                    <span className="text-lg font-semibold mb-2">
                      Access Token
                    </span>
                    <pre className="w-full h-full overflow-auto text-sm bg-gray-100 p-2 rounded">
                      {JSON.stringify(
                        session.tokens?.accessToken?.payload,
                        null,
                        2
                      )}
                    </pre>
                  </label>
                </div>
              </div>
            </fieldset>
          </>
        ) : null}
      </section>
      <section>
        <button
          onClick={() => query("queryWithGroup").then(setResult)}
          type="button"
        >
          query
        </button>
        <div>
          <pre>
            <code>{JSON.stringify(result, null, 2)}</code>
          </pre>
        </div>
      </section>
    </main>
  )
}

export function App() {
  return (
    <Authenticator>
      <Main />
    </Authenticator>
  )
}