Improve SSO auth flow

Use replaceState instead of a redirect to strip the loginToken
Put user into the same post-auth flows of E2ESetup
Skip UIA prompt in this post-auth flow, happy path is a server grace period

src/Lifecycle.ts

@@ -366,7 +366,7 @@ async function abortLogin() {
 //      The plan is to gradually move the localStorage access done here into
 //      SessionStore to avoid bugs where the view becomes out-of-sync with
 //      localStorage (e.g. isGuest etc.)
-async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }): Promise<boolean> {
+export async function restoreFromLocalStorage(opts?: { ignoreGuest?: boolean }): Promise<boolean> {
     const ignoreGuest = opts?.ignoreGuest;
 
     if (!localStorage) {

src/components/structures/MatrixChat.tsx

@@ -218,6 +218,7 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
     private screenAfterLogin?: IScreen;
     private windowWidth: number;
     private pageChanging: boolean;
+    private tokenLogin?: boolean;
     private accountPassword?: string;
     private accountPasswordTimer?: NodeJS.Timeout;
     private focusComposer: boolean;
@@ -323,13 +324,16 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
             Lifecycle.attemptTokenLogin(
                 this.props.realQueryParams,
                 this.props.defaultDeviceDisplayName,
-            ).then((loggedIn) => {
+            ).then(async (loggedIn) => {
                 if (loggedIn) {
+                    this.tokenLogin = true;
                     this.props.onTokenLoginCompleted();
 
-                    // don't do anything else until the page reloads - just stay in
+                    // Create and start the client
-                    // the 'loading' state.
+                    await Lifecycle.restoreFromLocalStorage({
-                    return;
+                        ignoreGuest: true,
+                    });
+                    return this.postLoginSetup();
                 }
 
                 // if the user has followed a login or register link, don't reanimate
@@ -353,6 +357,39 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
         CountlyAnalytics.instance.enable(/* anonymous = */ true);
     }
 
+    private async postLoginSetup() {
+        const cli = MatrixClientPeg.get();
+        const cryptoEnabled = cli.isCryptoEnabled();
+
+        const promisesList = [this.firstSyncPromise.promise];
+        if (cryptoEnabled) {
+            // wait for the client to finish downloading cross-signing keys for us so we
+            // know whether or not we have keys set up on this account
+            promisesList.push(cli.downloadKeys([cli.getUserId()]));
+        }
+
+        // Now update the state to say we're waiting for the first sync to complete rather
+        // than for the login to finish.
+        this.setState({ pendingInitialSync: true });
+
+        await Promise.all(promisesList);
+
+        if (!cryptoEnabled) {
+            this.setState({ pendingInitialSync: false });
+            return;
+        }
+
+        const crossSigningIsSetUp = cli.getStoredCrossSigningForUser(cli.getUserId());
+        if (crossSigningIsSetUp) {
+            this.setStateForNewView({ view: Views.COMPLETE_SECURITY });
+        } else if (await cli.doesServerSupportUnstableFeature("org.matrix.e2e_cross_signing")) {
+            this.setStateForNewView({ view: Views.E2E_SETUP });
+        } else {
+            this.onLoggedIn();
+        }
+        this.setState({ pendingInitialSync: false });
+    }
+
     // TODO: [REACT-WARNING] Replace with appropriate lifecycle stage
     // eslint-disable-next-line camelcase
     UNSAFE_componentWillUpdate(props, state) {
@@ -1833,40 +1870,7 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
 
         // Create and start the client
         await Lifecycle.setLoggedIn(credentials);
-
+        await this.postLoginSetup();
-        const cli = MatrixClientPeg.get();
-        const cryptoEnabled = cli.isCryptoEnabled();
-        if (!cryptoEnabled) {
-            this.onLoggedIn();
-        }
-
-        const promisesList = [this.firstSyncPromise.promise];
-        if (cryptoEnabled) {
-            // wait for the client to finish downloading cross-signing keys for us so we
-            // know whether or not we have keys set up on this account
-            promisesList.push(cli.downloadKeys([cli.getUserId()]));
-        }
-
-        // Now update the state to say we're waiting for the first sync to complete rather
-        // than for the login to finish.
-        this.setState({ pendingInitialSync: true });
-
-        await Promise.all(promisesList);
-
-        if (!cryptoEnabled) {
-            this.setState({ pendingInitialSync: false });
-            return;
-        }
-
-        const crossSigningIsSetUp = cli.getStoredCrossSigningForUser(cli.getUserId());
-        if (crossSigningIsSetUp) {
-            this.setStateForNewView({ view: Views.COMPLETE_SECURITY });
-        } else if (await cli.doesServerSupportUnstableFeature("org.matrix.e2e_cross_signing")) {
-            this.setStateForNewView({ view: Views.E2E_SETUP });
-        } else {
-            this.onLoggedIn();
-        }
-        this.setState({ pendingInitialSync: false });
     };
 
     // complete security / e2e setup has finished
@@ -1910,6 +1914,7 @@ export default class MatrixChat extends React.PureComponent<IProps, IState> {
                 <E2eSetup
                     onFinished={this.onCompleteSecurityE2eSetupFinished}
                     accountPassword={this.accountPassword}
+                    tokenLogin={!!this.tokenLogin}
                 />
             );
         } else if (this.state.view === Views.LOGGED_IN) {

src/components/structures/auth/E2eSetup.js

@@ -24,6 +24,7 @@ export default class E2eSetup extends React.Component {
     static propTypes = {
         onFinished: PropTypes.func.isRequired,
         accountPassword: PropTypes.string,
+        tokenLogin: PropTypes.bool,
     };
 
     render() {
@@ -33,6 +34,7 @@ export default class E2eSetup extends React.Component {
                     <CreateCrossSigningDialog
                         onFinished={this.props.onFinished}
                         accountPassword={this.props.accountPassword}
+                        tokenLogin={this.props.tokenLogin}
                     />
                 </CompleteSecurityBody>
             </AuthPage>

src/components/views/dialogs/security/CreateCrossSigningDialog.js

@@ -34,6 +34,7 @@ import InteractiveAuthDialog from '../InteractiveAuthDialog';
 export default class CreateCrossSigningDialog extends React.PureComponent {
     static propTypes = {
         accountPassword: PropTypes.string,
+        tokenLogin: PropTypes.bool,
     };
 
     constructor(props) {
@@ -96,6 +97,9 @@ export default class CreateCrossSigningDialog extends React.PureComponent {
                 user: MatrixClientPeg.get().getUserId(),
                 password: this.state.accountPassword,
             });
+        } else if (this.props.tokenLogin) {
+            // We are hoping the grace period is active
+            await makeRequest({});
         } else {
             const dialogAesthetics = {
                 [SSOAuthEntry.PHASE_PREAUTH]: {
@@ -144,6 +148,12 @@ export default class CreateCrossSigningDialog extends React.PureComponent {
             });
             this.props.onFinished(true);
         } catch (e) {
+            if (this.props.tokenLogin) {
+                // ignore any failures, we are relying on grace period here
+                this.props.onFinished();
+                return;
+            }
+
             this.setState({ error: e });
             console.error("Error bootstrapping cross-signing", e);
         }