From 82366d7cbebe5921cfb2ac950280dcb3207bf272 Mon Sep 17 00:00:00 2001
From: Aleksander <170264518+t-aleksander@users.noreply.github.com>
Date: Wed, 7 Oct 2026 13:16:12 +0200
Subject: [PATCH] update e2e tests

---
 e2e/tests/rateLimit.spec.ts | 10 +++++--
 e2e/utils/api/clientMfa.ts  | 57 +++++++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/e2e/tests/rateLimit.spec.ts b/e2e/tests/rateLimit.spec.ts
index d105d15ca..181e4c9d2 100644
--- a/e2e/tests/rateLimit.spec.ts
+++ b/e2e/tests/rateLimit.spec.ts
@@ -20,6 +20,7 @@ import {
   clientMfaStepStart,
   ClientProtocol,
   MfaMethod,
+  presharedKey,
   waitForProxy,
 } from '../utils/api/clientMfa';
 import { apiLogin } from '../utils/api/users';
@@ -52,7 +53,7 @@ const expectPreconditionError = async (response: APIResponse, message: string) =
 
 const expectConnected = async (response: APIResponse) => {
   expect(response.status()).toBe(200);
-  expect((await response.json()).preshared_key).toBeTruthy();
+  expect(await presharedKey(response)).toBeTruthy();
 };
 
 const countActivityEvents = async (
@@ -266,9 +267,12 @@ test.describe('Rate limiting', () => {
       method: MfaMethod.TOTP,
     }));
 
+    // The flow start initiates the first step, so it uses one request of the limit.
     let attempt = await clientMfaConnect(request, first);
-    for (let i = 1; i < VPN_INITIATE_LIMIT / 2; i++) {
-      attempt = await clientMfaConnect(request, first);
+    for (let i = 1; i < VPN_INITIATE_LIMIT; i++) {
+      expect(
+        (await clientMfaStepStart(request, attempt.token, first.method)).status(),
+      ).toBe(200);
     }
 
     await expectPreconditionError(
diff --git a/e2e/utils/api/clientMfa.ts b/e2e/utils/api/clientMfa.ts
index 627ec1309..5451b3e09 100644
--- a/e2e/utils/api/clientMfa.ts
+++ b/e2e/utils/api/clientMfa.ts
@@ -89,21 +89,28 @@ export const clientMfaStart = (
   api: APIRequestContext,
   device: ClientDevice,
 ): Promise<APIResponse> =>
-  api.post(proxyUrl('/client-mfa/start'), {
-    data: {
-      location_id: device.locationId,
-      pubkey: device.pubkey,
-      method: device.method,
-      selected_methods: device.protocol === 'legacy' ? [] : [device.method],
-    },
-  });
+  device.protocol === 'legacy'
+    ? api.post(proxyUrl('/client-mfa/start'), {
+        data: {
+          location_id: device.locationId,
+          pubkey: device.pubkey,
+          method: device.method,
+        },
+      })
+    : api.post(proxyUrl('/mfa-flow/start'), {
+        data: {
+          location_id: device.locationId,
+          pubkey: device.pubkey,
+          selected_methods: [device.method],
+        },
+      });
 
 export const clientMfaStepStart = (
   api: APIRequestContext,
   token: string,
   method: MfaMethod,
 ): Promise<APIResponse> =>
-  api.post(proxyUrl('/client-mfa/step-start'), { data: { token, method } });
+  api.post(proxyUrl('/mfa-flow/step-start'), { data: { token, method } });
 
 export const clientMfaConnect = async (
   api: APIRequestContext,
@@ -111,14 +118,14 @@ export const clientMfaConnect = async (
 ): Promise<MfaAttempt> => {
   const start = await clientMfaStart(api, device);
   expect(start.status()).toBe(200);
-  const { token } = await start.json();
-  expect(token).toBeTruthy();
+  const body = await start.json();
   if (device.protocol === 'legacy') {
-    return { token };
+    expect(body.token).toBeTruthy();
+    return { token: body.token };
   }
-  const step = await clientMfaStepStart(api, token, device.method);
-  expect(step.status()).toBe(200);
-  return { token, stepAttemptId: (await step.json()).step_attempt_id };
+  const accepted = body.outcome?.Accepted;
+  expect(accepted?.token).toBeTruthy();
+  return { token: accepted.token, stepAttemptId: accepted.first_step.step_attempt_id };
 };
 
 export const clientMfaFinish = (
@@ -126,6 +133,20 @@ export const clientMfaFinish = (
   attempt: MfaAttempt,
   code: string,
 ): Promise<APIResponse> =>
-  api.post(proxyUrl('/client-mfa/finish'), {
-    data: { token: attempt.token, code, step_attempt_id: attempt.stepAttemptId },
-  });
+  attempt.stepAttemptId === undefined
+    ? api.post(proxyUrl('/client-mfa/finish'), { data: { token: attempt.token, code } })
+    : api.post(proxyUrl('/mfa-flow/step-finish'), {
+        data: {
+          token: attempt.token,
+          step_attempt_id: attempt.stepAttemptId,
+          submission: { Code: { code } },
+        },
+      });
+
+// A multi-step client gets the key inside the completed step result.
+export const presharedKey = async (
+  response: APIResponse,
+): Promise<string | undefined> => {
+  const body = await response.json();
+  return body.preshared_key ?? body.result?.outcome?.Completed?.preshared_key;
+};