From 684c9285a6f06193a5ff1b7ef3421a2d7cd6f570 Mon Sep 17 00:00:00 2001
From: Ajay Bura <32841439+ajbura@users.noreply.github.com>
Date: Sun, 19 Jan 2025 21:30:13 +0530
Subject: [PATCH] add sso uia stage

---
 src/app/components/ActionUIA.tsx              | 17 +++-
 src/app/components/uia-stages/SSOStage.tsx    | 95 +++++++++++++++++++
 src/app/components/uia-stages/index.ts        |  1 +
 .../features/settings/sessions/Sessions.tsx   |  1 -
 4 files changed, 109 insertions(+), 5 deletions(-)
 create mode 100644 src/app/components/uia-stages/SSOStage.tsx

diff --git a/src/app/components/ActionUIA.tsx b/src/app/components/ActionUIA.tsx
index 88f0fe0e..b9cd122f 100644
--- a/src/app/components/ActionUIA.tsx
+++ b/src/app/components/ActionUIA.tsx
@@ -3,7 +3,8 @@ import { AuthDict, AuthType, IAuthData, UIAFlow } from 'matrix-js-sdk';
 import { getUIAFlowForStages } from '../utils/matrix-uia';
 import { useSupportedUIAFlows, useUIACompleted, useUIAFlow } from '../hooks/useUIAFlows';
 import { UIAFlowOverlay } from './UIAFlowOverlay';
-import { PasswordStage } from './uia-stages';
+import { PasswordStage, SSOStage } from './uia-stages';
+import { useMatrixClient } from '../hooks/useMatrixClient';
 
 export const SUPPORTED_IN_APP_UIA_STAGES = [AuthType.Password, AuthType.Sso];
 
@@ -14,13 +15,13 @@ export const pickUIAFlow = (uiaFlows: UIAFlow[]): UIAFlow | undefined => {
 };
 
 type ActionUIAProps = {
-  userId: string;
   authData: IAuthData;
   ongoingFlow: UIAFlow;
   action: (authDict: AuthDict) => void;
   onCancel: () => void;
 };
-export function ActionUIA({ userId, authData, ongoingFlow, action, onCancel }: ActionUIAProps) {
+export function ActionUIA({ authData, ongoingFlow, action, onCancel }: ActionUIAProps) {
+  const mx = useMatrixClient();
   const completed = useUIACompleted(authData);
   const { getStageToComplete } = useUIAFlow(authData, ongoingFlow);
 
@@ -35,7 +36,15 @@ export function ActionUIA({ userId, authData, ongoingFlow, action, onCancel }: A
     >
       {stageToComplete.type === AuthType.Password && (
         <PasswordStage
-          userId={userId}
+          userId={mx.getUserId()!}
+          stageData={stageToComplete}
+          onCancel={onCancel}
+          submitAuthDict={action}
+        />
+      )}
+      {stageToComplete.type === AuthType.Sso && stageToComplete.session && (
+        <SSOStage
+          ssoRedirectURL={mx.getFallbackAuthUrl(AuthType.Sso, stageToComplete.session)}
           stageData={stageToComplete}
           onCancel={onCancel}
           submitAuthDict={action}
diff --git a/src/app/components/uia-stages/SSOStage.tsx b/src/app/components/uia-stages/SSOStage.tsx
new file mode 100644
index 00000000..0c027e77
--- /dev/null
+++ b/src/app/components/uia-stages/SSOStage.tsx
@@ -0,0 +1,95 @@
+import { Box, Button, color, config, Dialog, Header, Icon, IconButton, Icons, Text } from 'folds';
+import React, { useCallback, useEffect, useState } from 'react';
+import { AuthType } from 'matrix-js-sdk';
+import { StageComponentProps } from './types';
+
+export function SSOStage({
+  ssoRedirectURL,
+  stageData,
+  submitAuthDict,
+  onCancel,
+}: StageComponentProps & {
+  ssoRedirectURL: string;
+}) {
+  const { errorCode, error, session } = stageData;
+  const [ssoWindow, setSSOWindow] = useState<Window>();
+
+  const handleSubmit = useCallback(
+    () =>
+      submitAuthDict({
+        type: AuthType.Sso,
+        session,
+      }),
+    [submitAuthDict, session]
+  );
+
+  const handleContinue = () => {
+    const w = window.open(ssoRedirectURL, '_blank');
+    setSSOWindow(w ?? undefined);
+  };
+
+  useEffect(() => {
+    const handleMessage = (evt: MessageEvent) => {
+      if (ssoWindow && evt.data === 'authDone' && evt.source === ssoWindow) {
+        ssoWindow.close();
+        setSSOWindow(undefined);
+        handleSubmit();
+      }
+    };
+
+    window.addEventListener('message', handleMessage);
+    return () => {
+      window.removeEventListener('message', handleMessage);
+    };
+  }, [ssoWindow, handleSubmit]);
+
+  return (
+    <Dialog>
+      <Header
+        style={{
+          padding: `0 ${config.space.S200} 0 ${config.space.S400}`,
+        }}
+        variant="Surface"
+        size="500"
+      >
+        <Box grow="Yes">
+          <Text size="H4">SSO Login</Text>
+        </Box>
+        <IconButton size="300" onClick={onCancel} radii="300">
+          <Icon src={Icons.Cross} />
+        </IconButton>
+      </Header>
+      <Box
+        style={{ padding: `0 ${config.space.S400} ${config.space.S400}` }}
+        direction="Column"
+        gap="400"
+      >
+        <Text size="T200">
+          To perform this action you need to authenticate yourself by SSO login.
+        </Text>
+        {errorCode && (
+          <Box alignItems="Center" gap="100" style={{ color: color.Critical.Main }}>
+            <Icon size="50" src={Icons.Warning} filled />
+            <Text size="T200">
+              <b>{`${errorCode}: ${error}`}</b>
+            </Text>
+          </Box>
+        )}
+
+        {ssoWindow ? (
+          <Button variant="Primary" onClick={handleSubmit}>
+            <Text as="span" size="B400">
+              Continue
+            </Text>
+          </Button>
+        ) : (
+          <Button variant="Primary" onClick={handleContinue}>
+            <Text as="span" size="B400">
+              Continue with SSO
+            </Text>
+          </Button>
+        )}
+      </Box>
+    </Dialog>
+  );
+}
diff --git a/src/app/components/uia-stages/index.ts b/src/app/components/uia-stages/index.ts
index 9b158989..c422d988 100644
--- a/src/app/components/uia-stages/index.ts
+++ b/src/app/components/uia-stages/index.ts
@@ -4,4 +4,5 @@ export * from './EmailStage';
 export * from './PasswordStage';
 export * from './ReCaptchaStage';
 export * from './RegistrationTokenStage';
+export * from './SSOStage';
 export * from './TermsStage';
diff --git a/src/app/features/settings/sessions/Sessions.tsx b/src/app/features/settings/sessions/Sessions.tsx
index 1521315e..dcc69091 100644
--- a/src/app/features/settings/sessions/Sessions.tsx
+++ b/src/app/features/settings/sessions/Sessions.tsx
@@ -404,7 +404,6 @@ export function Sessions({ requestClose }: SessionsProps) {
                         >
                           {(ongoingFlow) => (
                             <ActionUIA
-                              userId={mx.getUserId()!}
                               authData={authData}
                               ongoingFlow={ongoingFlow}
                               action={deleteDevices}