From 60bc4450f3c6d26bf4f2e9f17a4b6210f59d98d0 Mon Sep 17 00:00:00 2001
From: RMidhunSuresh <hi@midhun.dev>
Date: Wed, 16 Feb 2022 13:21:04 +0530
Subject: [PATCH] Use type from server

---
 src/matrix/registration/Registration.ts | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/matrix/registration/Registration.ts b/src/matrix/registration/Registration.ts
index 0d342639..8bfce8ad 100644
--- a/src/matrix/registration/Registration.ts
+++ b/src/matrix/registration/Registration.ts
@@ -18,7 +18,7 @@ import type {HomeServerApi} from "../net/HomeServerApi";
 import type {BaseRegistrationStage} from "./stages/BaseRegistrationStage";
 import {DummyAuth} from "./stages/DummyAuth";
 import {TermsAuth} from "./stages/TermsAuth";
-import {TokenAuth, TOKEN_AUTH_TYPE} from "./stages/TokenAuth";
+import {TokenAuth} from "./stages/TokenAuth";
 import type {
     AccountDetails,
     RegistrationFlow,
@@ -94,7 +94,9 @@ export class Registration {
                 this._sessionInfo = response;
                 return undefined;
             case 401:
-                if (response.completed?.includes(currentStage.type)) {
+                // Support unstable prefix for TokenAuth
+                const typeFromServer = (currentStage as TokenAuth).typeFromServer;
+                if (response.completed?.includes(typeFromServer ?? currentStage.type)) {
                     return currentStage.nextStage;
                 }
                 else {
@@ -109,8 +111,9 @@ export class Registration {
                 return new DummyAuth(session, params?.[type]);
             case "m.login.terms":
                 return new TermsAuth(session, params?.[type]);
-            case TOKEN_AUTH_TYPE:
-                return new TokenAuth(session, params?.[type]);
+            case "org.matrix.msc3231.login.registration_token":
+            case "m.login.registration_token":
+                return new TokenAuth(session, params?.[type], type);
             default:
                 throw new Error(`Unknown stage: ${type}`);
         }