From 10d356e2d6a4ac2d7c64125f31337abce0ce1979 Mon Sep 17 00:00:00 2001
From: Nick Wylynko <nick@wylynko.com>
Date: Sat, 9 Nov 2024 13:58:46 +0800
Subject: [PATCH] Pull supplied oidc authToken through to websocket request

---
 .../src/Providers/AWSWebSocketProvider/authHeaders.ts    | 9 ++++++++-
 packages/api-graphql/src/internals/events/index.ts       | 8 ++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/packages/api-graphql/src/Providers/AWSWebSocketProvider/authHeaders.ts b/packages/api-graphql/src/Providers/AWSWebSocketProvider/authHeaders.ts
index faa822fc55c..55eedb2866b 100644
--- a/packages/api-graphql/src/Providers/AWSWebSocketProvider/authHeaders.ts
+++ b/packages/api-graphql/src/Providers/AWSWebSocketProvider/authHeaders.ts
@@ -26,6 +26,13 @@ const awsAuthTokenHeader = async ({ host }: AWSAppSyncRealTimeAuthInput) => {
 	};
 };
 
+const oidcAuthTokenHeader = async ({ host, additionalCustomHeaders }: AWSAppSyncRealTimeAuthInput) => {
+	return {
+		Authorization: additionalCustomHeaders.Authorization,
+		host,
+	};
+};
+
 const awsRealTimeApiKeyHeader = async ({
 	apiKey,
 	host,
@@ -110,7 +117,7 @@ export const awsRealTimeHeaderBasedAuth = async ({
 	const headerHandler = {
 		apiKey: awsRealTimeApiKeyHeader,
 		iam: awsRealTimeIAMHeader,
-		oidc: awsAuthTokenHeader,
+		oidc: oidcAuthTokenHeader,
 		userPool: awsAuthTokenHeader,
 		lambda: customAuthHeader,
 		none: customAuthHeader,
diff --git a/packages/api-graphql/src/internals/events/index.ts b/packages/api-graphql/src/internals/events/index.ts
index c1101111fd4..3422a210d7e 100644
--- a/packages/api-graphql/src/internals/events/index.ts
+++ b/packages/api-graphql/src/internals/events/index.ts
@@ -48,6 +48,10 @@ async function connect(
 		providerOptions.authenticationType,
 	);
 
+	if (options?.authToken) {
+		providerOptions.authToken = options.authToken
+	}
+
 	await eventProvider.connect(providerOptions);
 
 	let _subscription: Subscription;
@@ -62,6 +66,10 @@ async function connect(
 			subscribeOptions.authenticationType,
 		);
 
+		if (subOptions?.authToken) {
+			subscribeOptions.authToken = subOptions.authToken
+		}
+
 		_subscription = eventProvider
 			.subscribe(subscribeOptions)
 			.subscribe(observer);