fix(appsync): appsync Event API integration assertion tests

packages/@aws-cdk-testing/framework-integ/package.json

@@ -35,6 +35,7 @@
     "@aws-sdk/client-acm": "3.632.0",
     "@aws-sdk/client-rds": "3.632.0",
     "@aws-sdk/client-s3": "3.632.0",
+    "@aws-sdk/client-cognito-identity-provider": "3.632.0",
     "axios": "1.7.8",
     "delay": "5.0.0"
   },

packages/@aws-cdk-testing/framework-integ/test/aws-appsync/test/integ-assets/appsync-js-channel-namespace-handler.js

@@ -1,3 +1,12 @@
+function enrichEvent(event) {
+  return {
+    id: event.id,
+    payload: {
+      ...event.payload,
+      newField: 'newField'
+    }
+  }
+}
 export function onPublish(ctx) {
-  return ctx.events.filter((event) => event.payload.odds > 0)
+  return ctx.events.map(enrichEvent);
 }

packages/@aws-cdk-testing/framework-integ/test/aws-appsync/test/integ-assets/eventapi-grant-assertion/index.js

@@ -4,6 +4,13 @@ import { HttpRequest } from '@smithy/protocol-http'
 import { SignatureV4 } from '@smithy/signature-v4'
 import { fromNodeProviderChain } from '@aws-sdk/credential-providers'
 import { Sha256 } from '@aws-crypto/sha256-js'
+import { 
+  CognitoIdentityProviderClient, 
+  SignUpCommand,
+  AdminConfirmSignUpCommand,
+  AdminDeleteUserCommand,
+  AdminInitiateAuthCommand,
+} from "@aws-sdk/client-cognito-identity-provider";
 
 // The default headers to to sign the request
 const DEFAULT_HEADERS = {
@@ -16,6 +23,56 @@ const AWS_APPSYNC_EVENTS_SUBPROTOCOL = 'aws-appsync-event-ws';
 const realtimeUrl = process.env.EVENT_API_REALTIME_URL;
 const httpUrl = process.env.EVENT_API_HTTP_URL;
 const region = process.env.AWS_REGION;
+const API_KEY = process.env.API_KEY;
+const USER_POOL_ID = process.env.USER_POOL_ID;
+const CLIENT_ID = process.env.CLIENT_ID;
+const { username, password } = generateUsernamePassword(12);
+
+const cognitoClient = new CognitoIdentityProviderClient();
+
+/**
+ * Utility function for generating a temporary password
+ * @param {int} length 
+ * @returns 
+ */
+function generateUsernamePassword(length) {
+  const uppercaseChars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ';
+  const lowercaseChars = 'abcdefghijklmnopqrstuvwxyz';
+  const numberChars = '0123456789';
+  const specialChars = '!@#$%&';
+  const allChars = uppercaseChars + lowercaseChars + numberChars + specialChars;
+
+  // Ensure length is at least 4 to accommodate required characters
+  const actualLength = Math.max(length, 4);
+
+  // Start with one character from each required set
+  let password = [
+    uppercaseChars.charAt(Math.floor(Math.random() * uppercaseChars.length)),
+    lowercaseChars.charAt(Math.floor(Math.random() * lowercaseChars.length)),
+    numberChars.charAt(Math.floor(Math.random() * numberChars.length)),
+    specialChars.charAt(Math.floor(Math.random() * specialChars.length))
+  ];
+
+  // Fill the rest with random characters
+  for (let i = 4; i < actualLength; i++) {
+    const randomIndex = Math.floor(Math.random() * allChars.length);
+    password.push(allChars.charAt(randomIndex));
+  }
+
+  // Shuffle the password array to randomize character positions
+  for (let i = password.length - 1; i > 0; i--) {
+    const j = Math.floor(Math.random() * (i + 1));
+    [password[i], password[j]] = [password[j], password[i]];
+  }
+
+  let username = '';
+  for (let i = 0; i < 6; i++) {
+    const randomIndex = Math.floor(Math.random() * lowercaseChars.length);
+    username += lowercaseChars.charAt(randomIndex);
+  }
+
+  return { username, password: password.join('') };
+}
 
 /**
  * Returns a signed authorization object
@@ -33,7 +90,7 @@ async function signWithAWSV4(httpDomain, region, body) {
     sha256: Sha256,
   })
 
-  const url = new URL(`https://${httpDomain}/event`)
+  const url = new URL(`${httpDomain}`)
   const request = new HttpRequest({
     method: 'POST',
     headers: {
@@ -55,13 +112,11 @@ async function signWithAWSV4(httpDomain, region, body) {
 
 /**
  * Returns a header value for the SubProtocol header
- * @param {string} httpDomain the AppSync Event API HTTP domain
- * @param {string} region the AWS region of your API
+ * @param {string} authHeaders the authorization headers
  * @returns string a header string
  */
-async function getAuthProtocolForIAM(httpDomain, region) {
-  const signed = await signWithAWSV4(httpDomain, region)
-  const based64UrlHeader = btoa(JSON.stringify(signed))
+function getAuthProtocolForIAM(authHeaders) {
+  const based64UrlHeader = btoa(JSON.stringify(authHeaders))
     .replace(/\+/g, '-') // Convert '+' to '-'
     .replace(/\//g, '_') // Convert '/' to '_'
     .replace(/=+$/, '') // Remove padding `=`
@@ -78,19 +133,107 @@ function sleep(ms) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
 
+/**
+ * Helper function for creating a Cognito user and confirming the user
+ * The function also deletes the user after the test is complete
+ * and it initiates and auth flow to get the ID token for testing the
+ * Event API auth flow with Cognito.
+ * @param {string} action - CREATE, DELETE, AUTH
+ * @returns 
+ */
+async function cognitoUserConfiguration(action) {
+  switch (action) {
+    case 'CREATE':
+      const signUpUserInput = {
+        ClientId: CLIENT_ID,
+        Username: username,
+        Password: password,
+      };
+      const signUpCommand = new SignUpCommand(signUpUserInput);
+      await cognitoClient.send(signUpCommand);
+      const confirmSignUpInput = {
+        UserPoolId: USER_POOL_ID,
+        Username: username,
+      };
+      const confirmSignUpCommand = new AdminConfirmSignUpCommand(confirmSignUpInput);
+      await cognitoClient.send(confirmSignUpCommand);
+      return {};
+    case 'DELETE':
+      const deleteUserInput = {
+        UserPoolId: USER_POOL_ID,
+        Username: username,
+      };
+      const deleteUserCommand = new AdminDeleteUserCommand(deleteUserInput);
+      await cognitoClient.send(deleteUserCommand);
+      return;
+    case 'AUTH':
+      const authInput = {
+        UserPoolId: USER_POOL_ID,
+        ClientId: CLIENT_ID,
+        AuthFlow: 'ADMIN_USER_PASSWORD_AUTH',
+        AuthParameters: {
+          USERNAME: username,
+          PASSWORD: password,
+        },
+      };
+      const authCommand = new AdminInitiateAuthCommand(authInput);
+      const authRes = await cognitoClient.send(authCommand);
+      return authRes.AuthenticationResult.IdToken;
+  }
+}
+
+/**
+ * Returns the appropriate headers depending on the auth mode selected
+ * @param {*} authMode - IAM, API_KEY, LAMBDA, USER_POOL, OIDC
+ * @param {*} event - the event payload for Publish operations, null by default
+ * @param {*} authToken - the token for LAMBDA auth modes
+ * @returns 
+ */
+async function getPublishAuthHeader(authMode, event={}, authToken='') {
+  const url = new URL(`${httpUrl}`)
+  const headers = {
+    host: url.hostname,
+  };
+
+  switch (authMode) {
+    case 'IAM':
+      return await signWithAWSV4(httpUrl, region, JSON.stringify(event));
+    case 'API_KEY':
+      return {
+        'x-api-key': `${API_KEY}`,
+        ...headers,
+      }
+    case 'USER_POOL':
+      return {
+        'Authorization': await cognitoUserConfiguration('AUTH'),
+        ...headers,
+      }
+    case 'LAMBDA':
+      return {
+        'Authorization': authToken,
+        ...headers,
+      }
+    default:
+      throw new Error(`Unknown auth mode ${authMode}`)
+  }
+}
+
 /**
  * Initiates a subscription to a channel and returns the response
  *
  * @param {string} channel the channel to subscribe to
+ * @param {string} authMode the authorization mode for the request
+ * @param {string} authToken the token used for Lambda auth mode
  * @param {boolean} triggerPub whether to also publish in the method
  * @returns {Object}
  */
-async function subscribe(channel, triggerPub=false) {
+async function subscribe(channel, authMode, authToken, triggerPub=false) {
   const response = {};
-  const auth = await getAuthProtocolForIAM(httpUrl, region)
+  const authHeader = await getPublishAuthHeader(authMode, {}, authToken);
+  const auth = getAuthProtocolForIAM(authHeader);
   const socket = await new Promise((resolve, reject) => {
     const socket = new WebSocket(
-      `wss://${realtimeUrl}/event/realtime`,
+      `${realtimeUrl}`,
       [AWS_APPSYNC_EVENTS_SUBPROTOCOL, auth],
       { headers: { ...DEFAULT_HEADERS } },
     )
@@ -138,12 +281,12 @@ async function subscribe(channel, triggerPub=false) {
     type: 'subscribe',
     id: crypto.randomUUID(),
     channel: subChannel,
-    authorization: await signWithAWSV4(httpUrl, region, JSON.stringify({ channel: subChannel })),
+    authorization: await getPublishAuthHeader(authMode, { channel: subChannel }, authToken),
   }));
 
   if (triggerPub) {
     await sleep(1000);
-    await publish(channel);
+    await publish(channel, authMode, authToken);
   }
   await sleep(3000);
   return response;
@@ -153,19 +296,21 @@ async function subscribe(channel, triggerPub=false) {
  * Publishes to a channel and returns the response
  *
  * @param {string} channel the channel to publish to
+ * @param {string} authMode the auth mode to use for publishing
+ * @param {string} authToken the auth token to use for Lambda auth mode
  * @returns {Object}
  */
-async function publish(channel) {
+async function publish(channel, authMode, authToken) {
   const event = {
     "channel": `/${channel}/test`,
     "events": [
       JSON.stringify({message:'Hello World!'})
     ]
   }
 
-  const response = await fetch(`https://${httpUrl}/event`, {
+  const response = await fetch(`${httpUrl}`, {
     method: 'POST',
-    headers: await signWithAWSV4(httpUrl, region, JSON.stringify(event)),
+    headers: await getPublishAuthHeader(authMode, event, authToken),
     body: JSON.stringify(event)
   });
 
@@ -190,18 +335,34 @@ async function publish(channel) {
 exports.handler = async function(event) {
   const pubSubAction = event.action;
   const channel = event.channel;
+  const authMode = event.authMode;
+  const authToken = event.authToken ?? '';
+  const isCustomEndpoint = event.customEndpoint ?? false;
 
+  // If custom endpoint, wait for 60 seconds for DNS to propagate
+  if (isCustomEndpoint) {
+    await sleep(60000);
+  }
+
+  if (authMode === 'USER_POOL') {
+    await cognitoUserConfiguration('CREATE');
+  }
+
+  let res;
   if (pubSubAction === 'publish') {
-    const res = await publish(channel);
+    res = await publish(channel, authMode, authToken);
     console.log(res);
-    return res;
   } else if (pubSubAction === 'subscribe') {
-    const res = await subscribe(channel, false);
+    res = await subscribe(channel, authMode, authToken, false);
     console.log(res);
-    return res;
   } else if (pubSubAction === 'pubSub') {
-    const res = await subscribe(channel, true);
+    res = await subscribe(channel, authMode, authToken, true);
     console.log(res);
-    return res;
   }
-};
+
+  if (authMode === 'USER_POOL') {
+    await cognitoUserConfiguration('DELETE');
+  }
+
+  return res;
+};