From 1761a4ec800624bef0b81875948cff50610324f8 Mon Sep 17 00:00:00 2001
From: Michael Telatynski <7t3chguy@gmail.com>
Date: Thu, 28 Jan 2021 10:45:29 +0000
Subject: [PATCH] Improve SSO UIA

Fixes Fallback UIA postmessage interface
Auto-closes SSO UIA tab when the user has completed the flow within it
Error for when auth stage is restarted because it failed
---
 src/components/structures/InteractiveAuth.js  |  8 ++-
 src/components/structures/MatrixChat.tsx      |  6 +-
 .../auth/InteractiveAuthEntryComponents.js    | 56 ++++++++++++++++---
 src/i18n/strings/en_EN.json                   |  1 +
 4 files changed, 62 insertions(+), 9 deletions(-)

diff --git a/src/components/structures/InteractiveAuth.js b/src/components/structures/InteractiveAuth.js
index c8fcd7e9ca..23e43da195 100644
--- a/src/components/structures/InteractiveAuth.js
+++ b/src/components/structures/InteractiveAuth.js
@@ -177,7 +177,13 @@ export default class InteractiveAuthComponent extends React.Component {
             stageState: stageState,
             errorText: stageState.error,
         }, () => {
-            if (oldStage != stageType) this._setFocus();
+            if (oldStage !== stageType) {
+                this._setFocus();
+            } else if (!stageState.error && this._stageComponent.current &&
+                this._stageComponent.current.attemptFailed
+            ) {
+                this._stageComponent.current.attemptFailed();
+            }
         });
     };
 
diff --git a/src/components/structures/MatrixChat.tsx b/src/components/structures/MatrixChat.tsx
index a399d92eaf..2f6c5bf353 100644
--- a/src/components/structures/MatrixChat.tsx
+++ b/src/components/structures/MatrixChat.tsx
@@ -325,9 +325,13 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
                 this.props.realQueryParams,
                 this.props.defaultDeviceDisplayName,
             ).then(async (loggedIn) => {
+                if (this.props.realQueryParams?.loginToken) {
+                    // remove the loginToken from the URL regardless
+                    this.props.onTokenLoginCompleted();
+                }
+
                 if (loggedIn) {
                     this.tokenLogin = true;
-                    this.props.onTokenLoginCompleted();
 
                     // Create and start the client
                     await Lifecycle.restoreFromLocalStorage({
diff --git a/src/components/views/auth/InteractiveAuthEntryComponents.js b/src/components/views/auth/InteractiveAuthEntryComponents.js
index 60e57afc98..c8469192d2 100644
--- a/src/components/views/auth/InteractiveAuthEntryComponents.js
+++ b/src/components/views/auth/InteractiveAuthEntryComponents.js
@@ -609,8 +609,12 @@ export class SSOAuthEntry extends React.Component {
             this.props.authSessionId,
         );
 
+        this._popupWindow = null;
+        window.addEventListener("message", this._onReceiveMessage);
+
         this.state = {
             phase: SSOAuthEntry.PHASE_PREAUTH,
+            attemptFailed: false,
         };
     }
 
@@ -618,12 +622,33 @@ export class SSOAuthEntry extends React.Component {
         this.props.onPhaseChange(SSOAuthEntry.PHASE_PREAUTH);
     }
 
+    componentWillUnmount() {
+        window.removeEventListener("message", this._onReceiveMessage);
+        if (this._popupWindow) {
+            this._popupWindow.close();
+        }
+    }
+
+    attemptFailed = () => {
+        this.setState({
+            attemptFailed: true,
+        });
+    };
+
+    _onReceiveMessage = event => {
+        if (event.data === "authDone" && event.origin === this.props.matrixClient.getHomeserverUrl()) {
+            if (this._popupWindow) {
+                this._popupWindow.close();
+            }
+        }
+    };
+
     onStartAuthClick = () => {
         // Note: We don't use PlatformPeg's startSsoAuth functions because we almost
         // certainly will need to open the thing in a new tab to avoid losing application
         // context.
 
-        window.open(this._ssoUrl, '_blank');
+        this._popupWindow = window.open(this._ssoUrl, "_blank");
         this.setState({phase: SSOAuthEntry.PHASE_POSTAUTH});
         this.props.onPhaseChange(SSOAuthEntry.PHASE_POSTAUTH);
     };
@@ -656,10 +681,28 @@ export class SSOAuthEntry extends React.Component {
             );
         }
 
-        return <div className='mx_InteractiveAuthEntryComponents_sso_buttons'>
-            {cancelButton}
-            {continueButton}
-        </div>;
+        let errorSection;
+        if (this.props.errorText) {
+            errorSection = (
+                <div className="error" role="alert">
+                    { this.props.errorText }
+                </div>
+            );
+        } else if (this.state.attemptFailed) {
+            errorSection = (
+                <div className="error" role="alert">
+                    { _t("Something went wrong in confirming your identity. Cancel and try again.") }
+                </div>
+            );
+        }
+
+        return <React.Fragment>
+            { errorSection }
+            <div className="mx_InteractiveAuthEntryComponents_sso_buttons">
+                {cancelButton}
+                {continueButton}
+            </div>
+        </React.Fragment>;
     }
 }
 
@@ -710,8 +753,7 @@ export class FallbackAuthEntry extends React.Component {
             this.props.loginType,
             this.props.authSessionId,
         );
-        this._popupWindow = window.open(url);
-        this._popupWindow.opener = null;
+        this._popupWindow = window.open(url, "_blank");
     };
 
     _onReceiveMessage = event => {
diff --git a/src/i18n/strings/en_EN.json b/src/i18n/strings/en_EN.json
index 8d047ea3f1..cc69a66d18 100644
--- a/src/i18n/strings/en_EN.json
+++ b/src/i18n/strings/en_EN.json
@@ -2334,6 +2334,7 @@
     "Please enter the code it contains:": "Please enter the code it contains:",
     "Code": "Code",
     "Submit": "Submit",
+    "Something went wrong in confirming your identity. Cancel and try again.": "Something went wrong in confirming your identity. Cancel and try again.",
     "Start authentication": "Start authentication",
     "Enter password": "Enter password",
     "Nice, strong password!": "Nice, strong password!",