From 3d515c0223e7a814f7c4d783f9e47ecfa6c25108 Mon Sep 17 00:00:00 2001
From: vgafxc <varadharaj.ganesan@scania.com>
Date: Sat, 14 Mar 2026 19:18:31 +0100
Subject: [PATCH] feat: add OIDC federation support (Entra ID / external IdP)

Add support for external OIDC identity providers (e.g., Microsoft Entra ID)
federated through Amazon Cognito. This enables three capabilities:

1. OAuth/OIDC redirect login: When OAuthDomain and OAuthRedirectUrl are
   present in aws-exports.json, Amplify is configured with loginWith.oauth
   so users are redirected to the external IdP instead of seeing the
   default Cognito login form.

2. Display federated user name: The top navigation bar now reads the
   user's name from the ID token payload (idToken.payload.name), falling
   back to username. For federated users, Cognito auto-generates a random
   username, so the ID token claim provides the actual display name.

3. Fix IoT policy attach for federated users: fetchUserAttributes() fails
   for OIDC-federated users, and the IoT policy attach was in the same
   try block after it, causing real-time metrics to silently break. Moved
   IoT policy attach before fetchUserAttributes and wrapped the latter in
   its own try/catch.

All changes are backward-compatible - when OAuthDomain is not configured,
behavior is identical to the current implementation.
---
 .../lambda/aws-exports-handler/index.ts          |  9 ++++++++-
 .../components/navigation/TopNavigationBar.tsx   | 16 +++++++++++++++-
 source/webui/src/contexts/UserContext.tsx        | 13 +++++++++----
 source/webui/src/main.tsx                        | 11 +++++++++++
 4 files changed, 43 insertions(+), 6 deletions(-)

diff --git a/source/infrastructure/lambda/aws-exports-handler/index.ts b/source/infrastructure/lambda/aws-exports-handler/index.ts
index 4b8217ed..aa8ade88 100644
--- a/source/infrastructure/lambda/aws-exports-handler/index.ts
+++ b/source/infrastructure/lambda/aws-exports-handler/index.ts
@@ -18,6 +18,8 @@ interface CloudFormationEvent {
     UserFilesBucketRegion: string;
     IoTEndpoint: string;
     IoTPolicy: string;
+    OAuthDomain?: string;
+    OAuthRedirectUrl?: string;
   };
   PhysicalResourceId?: string;
 }
@@ -44,11 +46,13 @@ export const handler = async (event: CloudFormationEvent): Promise<CloudFormatio
     UserFilesBucketRegion,
     IoTEndpoint,
     IoTPolicy,
+    OAuthDomain,
+    OAuthRedirectUrl,
   } = ResourceProperties;
 
   try {
     if (RequestType === "Create" || RequestType === "Update") {
-      const awsExports = {
+      const awsExports: Record<string, string> = {
         UserPoolId,
         PoolClientId,
         IdentityPoolId,
@@ -59,6 +63,9 @@ export const handler = async (event: CloudFormationEvent): Promise<CloudFormatio
         IoTPolicy,
       };
 
+      if (OAuthDomain) awsExports.OAuthDomain = OAuthDomain;
+      if (OAuthRedirectUrl) awsExports.OAuthRedirectUrl = OAuthRedirectUrl;
+
       await s3Client.send(
         new PutObjectCommand({
           Bucket: BucketName,
diff --git a/source/webui/src/components/navigation/TopNavigationBar.tsx b/source/webui/src/components/navigation/TopNavigationBar.tsx
index cf2d80fa..e8c82c83 100644
--- a/source/webui/src/components/navigation/TopNavigationBar.tsx
+++ b/source/webui/src/components/navigation/TopNavigationBar.tsx
@@ -3,9 +3,23 @@
 
 import { TopNavigation, TopNavigationProps } from "@cloudscape-design/components";
 import { useAuthenticator } from "@aws-amplify/ui-react";
+import { fetchAuthSession } from "aws-amplify/auth";
+import { useState, useEffect } from "react";
 
 export default function TopNavigationBar() {
   const { user, signOut } = useAuthenticator();
+  const [displayName, setDisplayName] = useState<string | null>(null);
+
+  useEffect(() => {
+    fetchAuthSession()
+      .then((session) => {
+        const payload = session.tokens?.idToken?.payload;
+        if (payload) {
+          setDisplayName((payload.name as string) || (payload.email as string) || null);
+        }
+      })
+      .catch(() => {});
+  }, [user]);
 
   const solutionIdentity: TopNavigationProps.Identity = {
     href: "/",
@@ -20,7 +34,7 @@ export default function TopNavigationBar() {
   const utilities: TopNavigationProps.Utility[] = [
     {
       type: "menu-dropdown",
-      text: user.username ?? "User",
+      text: displayName ?? user.username ?? "User",
       iconName: "user-profile",
       items: [
         {
diff --git a/source/webui/src/contexts/UserContext.tsx b/source/webui/src/contexts/UserContext.tsx
index eba752ef..778b9a15 100644
--- a/source/webui/src/contexts/UserContext.tsx
+++ b/source/webui/src/contexts/UserContext.tsx
@@ -58,15 +58,20 @@ export const UserContextProvider = (props: { children: ReactNode }) => {
         ...responseUser,
       });
       try {
-        const userAttributesOutput = await fetchUserAttributes();
-        setEmail(userAttributesOutput.email ?? null);
-
-        // Attach IoT policy when user is authenticated
+        // Attach IoT policy first — fetchUserAttributes may fail for federated users
+        // and must not block IoT setup needed for real-time metrics
         const response = await fetch("/aws-exports.json");
         const config = await response.json();
         if (config.IoTPolicy) {
           await attachIoTPolicy(config.IoTPolicy);
         }
+
+        try {
+          const userAttributesOutput = await fetchUserAttributes();
+          setEmail(userAttributesOutput.email ?? null);
+        } catch (attrError) {
+          console.log("fetchUserAttributes failed (expected for federated users):", attrError);
+        }
       } catch (e) {
         console.log(e);
       }
diff --git a/source/webui/src/main.tsx b/source/webui/src/main.tsx
index fcf4db7f..0abb0eea 100644
--- a/source/webui/src/main.tsx
+++ b/source/webui/src/main.tsx
@@ -41,6 +41,17 @@ getRuntimeConfig().then((json) => {
         userPoolId: json.UserPoolId,
         userPoolClientId: json.PoolClientId,
         identityPoolId: json.IdentityPoolId,
+        ...(json.OAuthDomain && {
+          loginWith: {
+            oauth: {
+              domain: json.OAuthDomain,
+              scopes: ["email", "openid", "profile"],
+              redirectSignIn: [json.OAuthRedirectUrl],
+              redirectSignOut: [json.OAuthRedirectUrl],
+              responseType: "code",
+            },
+          },
+        }),
       },
     },
     API: {