From 4d9e089643b5f9faf4a6892f08ec757d2625144f Mon Sep 17 00:00:00 2001
From: alina sireneva <alina@tei.su>
Date: Sun, 30 Jun 2024 23:07:11 +0300
Subject: [PATCH] fix(core): gracefully handle wrong password in qr login

---
 packages/core/src/highlevel/client.ts         |  8 ++++
 .../src/highlevel/methods/auth/sign-in-qr.ts  | 41 ++++++++++++++++++-
 .../core/src/highlevel/methods/auth/start.ts  |  3 ++
 3 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/packages/core/src/highlevel/client.ts b/packages/core/src/highlevel/client.ts
index 73df7568..bd70098b 100644
--- a/packages/core/src/highlevel/client.ts
+++ b/packages/core/src/highlevel/client.ts
@@ -732,6 +732,14 @@ export interface TelegramClient extends ITelegramClient {
         /** Password for 2FA */
         password?: MaybeDynamic<string>
 
+        /**
+         * Function that will be called after the server has rejected the password.
+         *
+         * Note that in case {@link password} is not a function,
+         * this callback will never be called, and an error will be thrown instead.
+         */
+        invalidPasswordCallback?: () => MaybePromise<void>
+
         /** Abort signal */
         abortSignal?: AbortSignal
     }): Promise<User>
diff --git a/packages/core/src/highlevel/methods/auth/sign-in-qr.ts b/packages/core/src/highlevel/methods/auth/sign-in-qr.ts
index 5245b55c..07ef3498 100644
--- a/packages/core/src/highlevel/methods/auth/sign-in-qr.ts
+++ b/packages/core/src/highlevel/methods/auth/sign-in-qr.ts
@@ -1,6 +1,7 @@
 import { tl } from '@mtcute/tl'
 
 import { getPlatform } from '../../../platform.js'
+import { MaybePromise } from '../../../types/utils.js'
 import { ControllablePromise, createControllablePromise } from '../../../utils/controllable-promise.js'
 import { sleepWithAbort } from '../../../utils/misc-utils.js'
 import { assertTypeIs } from '../../../utils/type-assertions.js'
@@ -29,6 +30,14 @@ export async function signInQr(
         /** Password for 2FA */
         password?: MaybeDynamic<string>
 
+        /**
+         * Function that will be called after the server has rejected the password.
+         *
+         * Note that in case {@link password} is not a function,
+         * this callback will never be called, and an error will be thrown instead.
+         */
+        invalidPasswordCallback?: () => MaybePromise<void>
+
         /** Abort signal */
         abortSignal?: AbortSignal
     },
@@ -59,6 +68,34 @@ export async function signInQr(
         waiter?.reject(abortSignal.reason)
     })
 
+    async function handle2fa(input: MaybeDynamic<string>) {
+        const isDynamic = typeof input === 'function'
+
+        while (true) {
+            const password = await resolveMaybeDynamic(input)
+
+            try {
+                return await checkPassword(client, password)
+            } catch (e) {
+                if (tl.RpcError.is(e, 'PASSWORD_HASH_INVALID')) {
+                    if (!isDynamic) {
+                        throw e
+                    }
+
+                    if (params.invalidPasswordCallback) {
+                        await params.invalidPasswordCallback?.()
+                    } else {
+                        // eslint-disable-next-line no-console
+                        console.log('Invalid password. Please try again')
+                    }
+                    continue
+                }
+
+                throw e
+            }
+        }
+    }
+
     try {
         const { id, hash } = await client.getApiCrenetials()
         const platform = getPlatform()
@@ -78,7 +115,7 @@ export async function signInQr(
                 )
             } catch (e) {
                 if (tl.RpcError.is(e, 'SESSION_PASSWORD_NEEDED') && params.password) {
-                    return checkPassword(client, await resolveMaybeDynamic(params.password))
+                    return handle2fa(params.password)
                 }
 
                 throw e
@@ -109,7 +146,7 @@ export async function signInQr(
                         )
                     } catch (e) {
                         if (tl.RpcError.is(e, 'SESSION_PASSWORD_NEEDED') && params.password) {
-                            return checkPassword(client, await resolveMaybeDynamic(params.password))
+                            return handle2fa(params.password)
                         }
 
                         throw e
diff --git a/packages/core/src/highlevel/methods/auth/start.ts b/packages/core/src/highlevel/methods/auth/start.ts
index 2b0cb0b6..be856665 100644
--- a/packages/core/src/highlevel/methods/auth/start.ts
+++ b/packages/core/src/highlevel/methods/auth/start.ts
@@ -274,6 +274,9 @@ export async function start(
         return await signInQr(client, {
             onUrlUpdated: params.qrCodeHandler,
             password: params.password,
+            invalidPasswordCallback: params.invalidCodeCallback ?
+                () => params.invalidCodeCallback!('password') :
+                undefined,
             abortSignal,
         })
     }