From f46d2c1bf5ee01432bc2ef09296d72c194534ecf Mon Sep 17 00:00:00 2001
From: Bruno Windels <274386+bwindels@users.noreply.github.com>
Date: Fri, 14 Oct 2022 10:17:32 +0200
Subject: [PATCH] make startWithFinishedRegistration more broadly useful

---
 src/matrix/Client.js                    | 25 ++++++++++---------------
 src/matrix/registration/Registration.ts | 16 ++++++++++++----
 src/matrix/registration/types.ts        |  7 +++++++
 3 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/src/matrix/Client.js b/src/matrix/Client.js
index 2670ce07..4d6f7bb5 100644
--- a/src/matrix/Client.js
+++ b/src/matrix/Client.js
@@ -146,18 +146,13 @@ export class Client {
         return registration;
     }
 
-    async startWithFinishedRegistration(registration) {
-        this._platform.logger.run("startWithFinishedRegistration", async (log) => {
-            const sessionInfo = registration.sessionInfo;
-            if (!sessionInfo) {
-                throw new Error("Registration.sessionInfo is not available; are you sure that registration is finished?");
-            }
-            await this.startWithSessionInfo({
-                accessToken: sessionInfo.access_token,
-                deviceId: sessionInfo.device_id,
-                userId: sessionInfo.user_id,
-                homeserver: registration.homeserver,
-            }, true, log);
+    /** Method to start client after registration or with given access token.
+     * To start the client after registering, use `startWithAuthData(registration.authData)`.
+     * `homeserver` won't be resolved or normalized using this method,
+     * use `lookupHomeserver` first if needed (not needed after registration) */
+    async startWithAuthData({accessToken, deviceId, userId, homeserver}) {
+        this._platform.logger.run("startWithAuthData", async (log) => {
+            await this._createSessionAfterAuth({accessToken, deviceId, userId, homeserver}, true, log);
         });
     }
 
@@ -200,12 +195,12 @@ export class Client {
                 }
                 return;
             }
-            await this.startWithSessionInfo(sessionInfo, inspectAccountSetup, log);
+            await this._createSessionAfterAuth(sessionInfo, inspectAccountSetup, log);
         });
     }
 
-    async startWithSessionInfo({deviceId, userId, accessToken, homeserver}, inspectAccountSetup, log) {
-        await log.wrap("startWithSessionInfo", async (l) => {
+    async _createSessionAfterAuth({deviceId, userId, accessToken, homeserver}, inspectAccountSetup, log) {
+        await log.wrap("_createSessionAfterAuth", async (l) => {
             const id = this.createNewSessionId();
             const lastUsed = this._platform.clock.now();
             const sessionInfo = {
diff --git a/src/matrix/registration/Registration.ts b/src/matrix/registration/Registration.ts
index 82e3ade5..5440d66d 100644
--- a/src/matrix/registration/Registration.ts
+++ b/src/matrix/registration/Registration.ts
@@ -25,6 +25,7 @@ import type {
     RegistrationResponseMoreDataNeeded,
     RegistrationResponse,
     RegistrationResponseSuccess,
+    AuthData,
     RegistrationParams,
 } from "./types";
 
@@ -34,7 +35,7 @@ export class Registration {
     private readonly _hsApi: HomeServerApi;
     private readonly _accountDetails: AccountDetails;
     private readonly _flowSelector: FlowSelector;
-    private _sessionInfo?: RegistrationResponseSuccess;
+    private _registerResponse?: RegistrationResponseSuccess;
     public readonly homeserver: string;
 
     constructor(homeserver: string, hsApi: HomeServerApi, accountDetails: AccountDetails, flowSelector?: FlowSelector) {
@@ -93,7 +94,7 @@ export class Registration {
     private async parseRegistrationResponse(response: RegistrationResponse, currentStage: BaseRegistrationStage) {
         switch (response.status) {
             case 200:
-                this._sessionInfo = response;
+                this._registerResponse = response;
                 return undefined;
             case 401:
                 if (response.completed?.includes(currentStage.type)) {
@@ -119,7 +120,14 @@ export class Registration {
         }
     }
 
-    get sessionInfo(): RegistrationResponseSuccess | undefined {
-        return this._sessionInfo;
+    get authData(): AuthData | undefined {
+        if (this._registerResponse) {
+            return {
+                accessToken: this._registerResponse.access_token,
+                homeserver: this.homeserver,
+                userId: this._registerResponse.user_id,
+                deviceId: this._registerResponse.device_id,
+            };
+        }
     }
 }
diff --git a/src/matrix/registration/types.ts b/src/matrix/registration/types.ts
index f1ddbe98..7e0d01ff 100644
--- a/src/matrix/registration/types.ts
+++ b/src/matrix/registration/types.ts
@@ -38,6 +38,13 @@ export type RegistrationResponseSuccess = {
     status: 200;
 }
 
+export type AuthData = {
+    userId: string;
+    deviceId: string;
+    homeserver: string;
+    accessToken?: string;
+}
+
 export type RegistrationFlow = {
     stages: string[];
 }