Prioritise entirely supported flows for UIA (#3402)

* Prioritise entirely supported flows for UIA

* Add tests

* Fix types

* Apply suggestions from code review

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

* Update src/interactive-auth.ts

---------

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>

spec/unit/interactive-auth.spec.ts

@@ -517,4 +517,47 @@ describe("InteractiveAuth", () => {
             expect(ia.getEmailSid()).toEqual(sid);
         });
     });
+
+    it("should prioritise shorter flows", async () => {
+        const doRequest = jest.fn();
+        const stateUpdated = jest.fn();
+
+        const ia = new InteractiveAuth({
+            matrixClient: getFakeClient(),
+            doRequest: doRequest,
+            stateUpdated: stateUpdated,
+            requestEmailToken: jest.fn(),
+            authData: {
+                session: "sessionId",
+                flows: [{ stages: [AuthType.Recaptcha, AuthType.Password] }, { stages: [AuthType.Password] }],
+                params: {},
+            },
+        });
+
+        // @ts-ignore
+        ia.chooseStage();
+        expect(ia.getChosenFlow()?.stages).toEqual([AuthType.Password]);
+    });
+
+    it("should prioritise flows with entirely supported stages", async () => {
+        const doRequest = jest.fn();
+        const stateUpdated = jest.fn();
+
+        const ia = new InteractiveAuth({
+            matrixClient: getFakeClient(),
+            doRequest: doRequest,
+            stateUpdated: stateUpdated,
+            requestEmailToken: jest.fn(),
+            authData: {
+                session: "sessionId",
+                flows: [{ stages: ["com.devture.shared_secret_auth"] }, { stages: [AuthType.Password] }],
+                params: {},
+            },
+            supportedStages: [AuthType.Password],
+        });
+
+        // @ts-ignore
+        ia.chooseStage();
+        expect(ia.getChosenFlow()?.stages).toEqual([AuthType.Password]);
+    });
 });

src/interactive-auth.ts

@@ -26,7 +26,7 @@ const EMAIL_STAGE_TYPE = "m.login.email.identity";
 const MSISDN_STAGE_TYPE = "m.login.msisdn";
 
 export interface UIAFlow {
-    stages: AuthType[];
+    stages: Array<AuthType | string>;
 }
 
 export interface IInputs {
@@ -156,6 +156,14 @@ interface IOpts {
      */
     emailSid?: string;
 
+    /**
+     * If specified, will prefer flows which entirely consist of listed stages.
+     * These should normally be of type AuthTypes but can be string when supporting custom auth stages.
+     *
+     * This can be used to avoid needing the fallback mechanism.
+     */
+    supportedStages?: Array<AuthType | string>;
+
     /**
      * Called with the new auth dict to submit the request.
      * Also passes a second deprecated arg which is a flag set to true if this request is a background request.
@@ -176,7 +184,7 @@ interface IOpts {
      *     m.login.email.identity:
      *         * emailSid: string, the sid of the active email auth session
      */
-    stateUpdated(nextStage: AuthType, status: IStageStatus): void;
+    stateUpdated(nextStage: AuthType | string, status: IStageStatus): void;
 
     /**
      * A function that takes the email address (string), clientSecret (string), attempt number (int) and
@@ -216,6 +224,7 @@ export class InteractiveAuth {
     private readonly busyChangedCallback?: IOpts["busyChanged"];
     private readonly stateUpdatedCallback: IOpts["stateUpdated"];
     private readonly requestEmailTokenCallback: IOpts["requestEmailToken"];
+    private readonly supportedStages?: Set<string>;
 
     private data: IAuthData;
     private emailSid?: string;
@@ -243,6 +252,7 @@ export class InteractiveAuth {
         if (opts.sessionId) this.data.session = opts.sessionId;
         this.clientSecret = opts.clientSecret || this.matrixClient.generateClientSecret();
         this.emailSid = opts.emailSid;
+        if (opts.supportedStages !== undefined) this.supportedStages = new Set(opts.supportedStages);
     }
 
     /**
@@ -571,7 +581,7 @@ export class InteractiveAuth {
      * @returns login type
      * @throws {@link NoAuthFlowFoundError} If no suitable authentication flow can be found
      */
-    private chooseStage(): AuthType | undefined {
+    private chooseStage(): AuthType | string | undefined {
         if (this.chosenFlow === null) {
             this.chosenFlow = this.chooseFlow();
         }
@@ -581,6 +591,17 @@ export class InteractiveAuth {
         return nextStage;
     }
 
+    // Returns a low number for flows we consider best. Counts increase for longer flows and even more so
+    // for flows which contain stages not listed in `supportedStages`.
+    private scoreFlow(flow: UIAFlow): number {
+        let score = flow.stages.length;
+        if (this.supportedStages !== undefined) {
+            // Add 10 points to the score for each unsupported stage in the flow.
+            score += flow.stages.filter((stage) => !this.supportedStages!.has(stage)).length * 10;
+        }
+        return score;
+    }
+
     /**
      * Pick one of the flows from the returned list
      * If a flow using all of the inputs is found, it will
@@ -603,6 +624,10 @@ export class InteractiveAuth {
         const haveEmail = Boolean(this.inputs.emailAddress) || Boolean(this.emailSid);
         const haveMsisdn = Boolean(this.inputs.phoneCountry) && Boolean(this.inputs.phoneNumber);
 
+        // Flows are not represented in a significant order, so we can choose any we support best
+        // Sort flows based on how many unsupported stages they contain ascending
+        flows.sort((a, b) => this.scoreFlow(a) - this.scoreFlow(b));
+
         for (const flow of flows) {
             let flowHasEmail = false;
             let flowHasMsisdn = false;
@@ -633,7 +658,7 @@ export class InteractiveAuth {
      * @internal
      * @returns login type
      */
-    private firstUncompletedStage(flow: UIAFlow): AuthType | undefined {
+    private firstUncompletedStage(flow: UIAFlow): AuthType | string | undefined {
         const completed = this.data.completed || [];
         return flow.stages.find((stageType) => !completed.includes(stageType));
     }