From 1b96d049c731ac448635e4daac597b1bc6a37648 Mon Sep 17 00:00:00 2001
From: Verticool <darkeroman@gmail.com>
Date: Tue, 18 Mar 2025 22:42:25 +0600
Subject: [PATCH] useAuth additional validation checks: localStorage ton_proof
 save > use > remove & selectWallet check

---
 src/pages/view-content/index.tsx  |   2 +-
 src/shared/services/auth/index.ts | 358 ++++++++++++++++++++++--------
 2 files changed, 270 insertions(+), 90 deletions(-)

diff --git a/src/pages/view-content/index.tsx b/src/pages/view-content/index.tsx
index 2d8c088..0163501 100644
--- a/src/pages/view-content/index.tsx
+++ b/src/pages/view-content/index.tsx
@@ -232,7 +232,7 @@ export const ViewContentPage = () => {
                 {content?.data?.downloadable && (
                     <Button
                         onClick={() => handleDwnldContent()}
-                        className={'h-[48px] bg-darkred mb-4'}
+                        className={'h-[48px] mb-4'}
                         label={`Скачать контент`}
                     />
                 )}
diff --git a/src/shared/services/auth/index.ts b/src/shared/services/auth/index.ts
index f4197a4..aed764d 100644
--- a/src/shared/services/auth/index.ts
+++ b/src/shared/services/auth/index.ts
@@ -1,34 +1,36 @@
-import { useRef } from 'react';
+import { useRef, useEffect } from 'react';
 import { useTonConnectUI } from '@tonconnect/ui-react';
 import { useMutation } from 'react-query';
 import { request } from '~/shared/libs';
 import { useWebApp } from '@vkruglikov/react-telegram-web-app';
 
 const sessionStorageKey = 'auth_v1_token';
+const tonProofStorageKey = 'stored_ton_proof';
 const payloadTTLMS = 1000 * 60 * 20;
 
 export const useAuth = () => {
     const WebApp = useWebApp();
-    // const wallet = useTonWallet();
     const [tonConnectUI] = useTonConnectUI();
     const interval = useRef<ReturnType<typeof setInterval> | undefined>();
 
-    const waitForWalletProof = async () => {
-        return new Promise((resolve, reject) => {
-            const timeout = setTimeout(() => reject(new Error('Timeout waiting for proof')), 30000);
-            const checkProof = setInterval(() => {
-                const currentWallet = tonConnectUI.wallet;
-                if (
-                    currentWallet?.connectItems?.tonProof &&
-                    !('error' in currentWallet.connectItems.tonProof)
-                ) {
-                    clearInterval(checkProof);
-                    clearTimeout(timeout);
-                    resolve(currentWallet.connectItems.tonProof.proof);
-                }
-            }, 500);
-        });
-    };
+    // Store ton_proof when it becomes available
+    useEffect(() => {
+        if (
+            tonConnectUI.wallet?.connectItems?.tonProof &&
+            !('error' in tonConnectUI.wallet.connectItems.tonProof) &&
+            tonConnectUI.wallet.account
+        ) {
+            console.log('DEBUG: Storing ton_proof for future use');
+            localStorage.setItem(
+                tonProofStorageKey,
+                JSON.stringify({
+                    timestamp: Date.now(),
+                    account: tonConnectUI.wallet.account,
+                    proof: tonConnectUI.wallet.connectItems.tonProof.proof,
+                })
+            );
+        }
+    }, [tonConnectUI.wallet?.connectItems?.tonProof, tonConnectUI.wallet?.account]);
 
     const makeAuthRequest = async (params: {
         twa_data: string;
@@ -37,26 +39,94 @@ export const useAuth = () => {
             ton_proof: any;
         };
     }) => {
-        const res = await request.post<{
-            connected_wallet: null | {
-                version: string;
-                address: string;
-                ton_balance: string;
-            };
-            auth_v1_token: string;
-        }>('/auth.twa', params);
+        try {
+            const res = await request.post<{
+                connected_wallet: null | {
+                    version: string;
+                    address: string;
+                    ton_balance: string;
+                };
+                auth_v1_token: string;
+            }>('/auth.twa', params);
 
-        if (res?.data?.auth_v1_token) {
-            localStorage.setItem(sessionStorageKey, res.data.auth_v1_token);
-        } else {
-            throw new Error('Failed to get auth token');
+            if (res?.data?.auth_v1_token) {
+                localStorage.setItem(sessionStorageKey, res.data.auth_v1_token);
+
+                // If we sent a proof, it was accepted, so keep record of that
+                if (params.ton_proof) {
+                    console.log('DEBUG: Auth with proof successful');
+                }
+            } else {
+                throw new Error('Failed to get auth token');
+            }
+            return res;
+        } catch (error) {
+            // If we were using ton_proof and it failed, clear stored proof
+            if (params.ton_proof) {
+                console.log('DEBUG: Auth with proof failed, clearing stored proof');
+                localStorage.removeItem(tonProofStorageKey);
+            }
+            throw error;
         }
-        return res;
     };
 
     const makeSelectWalletRequest = async (params: { wallet_address: string }) => {
-        const res = await request.post('/auth.selectWallet', params);
-        return res;
+        try {
+            const res = await request.post('/auth.selectWallet', params);
+            return res;
+        } catch (error: any) {
+            // Check for 404 error (wallet not found or invalid)
+            if (error.response?.status === 404) {
+                console.log('DEBUG: Wallet selection failed with 404, disconnecting');
+                await tonConnectUI.disconnect();
+                localStorage.removeItem(sessionStorageKey);
+                localStorage.removeItem(tonProofStorageKey);
+            }
+            throw error;
+        }
+    };
+
+    // Helper to prepare the connection parameters with proof requirements
+    const prepareConnectParams = async () => {
+        console.log('DEBUG: Preparing connect parameters');
+
+        // Set to loading state first
+        tonConnectUI.setConnectRequestParameters({ state: 'loading' });
+
+        try {
+            // Get the payload/token from backend
+            const value = await request.post<{ auth_v1_token: string }>('/auth.twa', {
+                twa_data: WebApp.initData,
+            });
+
+            if (value?.data?.auth_v1_token) {
+                console.log('DEBUG: Got token for connect params');
+
+                // Set the parameters to ready with tonProof requirement
+                tonConnectUI.setConnectRequestParameters({
+                    state: 'ready',
+                    value: { tonProof: value.data.auth_v1_token },
+                });
+                return true;
+            } else {
+                console.log('DEBUG: No token received for connect params');
+                tonConnectUI.setConnectRequestParameters(null);
+                return false;
+            }
+        } catch (error) {
+            console.error('DEBUG: Error preparing connect params:', error);
+            tonConnectUI.setConnectRequestParameters(null);
+            return false;
+        }
+    };
+
+    // Helper to check if connection is capable of transactions
+    const isConnectionValid = () => {
+        if (!tonConnectUI.connected || !tonConnectUI.wallet) {
+            return false;
+        }
+
+        return true;
     };
 
     return useMutation(['auth'], async () => {
@@ -64,68 +134,178 @@ export const useAuth = () => {
         let authResult;
         console.log('DEBUG: Starting auth flow');
 
-        // Case 1: Not connected - need to connect and get proof
-        if (!tonConnectUI.connected) {
-            console.log('DEBUG: No wallet connection, starting flow');
-            localStorage.removeItem(sessionStorageKey);
+        try {
+            // Case 1: Not connected - need to connect and get proof
+            if (!tonConnectUI.connected) {
+                console.log('DEBUG: No wallet connection, starting flow');
+                localStorage.removeItem(sessionStorageKey);
 
-            const refreshPayload = async () => {
-                tonConnectUI.setConnectRequestParameters({ state: 'loading' });
-                const value = await request.post<{ auth_v1_token: string }>('/auth.twa', {
-                    twa_data: WebApp.initData,
-                });
+                // Prepare connection parameters (this sets up the proof requirement)
+                const prepared = await prepareConnectParams();
 
-                if (value?.data?.auth_v1_token) {
-                    tonConnectUI.setConnectRequestParameters({
-                        state: 'ready',
-                        value: { tonProof: value.data.auth_v1_token },
-                    });
-                } else {
-                    tonConnectUI.setConnectRequestParameters(null);
+                if (!prepared) {
+                    throw new Error('Failed to prepare connection parameters');
                 }
-            };
 
-            await refreshPayload();
-            interval.current = setInterval(refreshPayload, payloadTTLMS);
+                // Start periodic refresh of the payload
+                interval.current = setInterval(prepareConnectParams, payloadTTLMS);
 
-            const tonProof = await waitForWalletProof();
-            console.log('DEBUG: Got initial proof', tonProof);
+                // Open the modal - this will not resolve until connection or cancellation
+                await tonConnectUI.openModal();
 
-            authResult = await makeAuthRequest({
-                twa_data: WebApp.initData,
-                ton_proof: {
-                    account: tonConnectUI.wallet!.account,
-                    ton_proof: tonProof,
-                },
-            });
-        } else {
-            // Case 3: Connected without proof - already authenticated
-            console.log('DEBUG: Connected without proof, proceeding without it');
-            authResult = await makeAuthRequest({
-                twa_data: WebApp.initData,
-            });
+                // Check if connection was successful
+                if (!tonConnectUI.connected) {
+                    console.log('DEBUG: Connection cancelled or failed');
+                    throw new Error('Wallet connection cancelled or failed');
+                }
+
+                // Check if we have a proof after connection
+                if (
+                    tonConnectUI.wallet?.connectItems?.tonProof &&
+                    !('error' in tonConnectUI.wallet.connectItems.tonProof)
+                ) {
+                    console.log('DEBUG: Got proof after connection');
+
+                    try {
+                        // Try auth with the proof
+                        authResult = await makeAuthRequest({
+                            twa_data: WebApp.initData,
+                            ton_proof: {
+                                account: tonConnectUI.wallet.account,
+                                ton_proof: tonConnectUI.wallet.connectItems.tonProof.proof,
+                            },
+                        });
+                    } catch (error) {
+                        // If auth with proof fails, throw the error
+                        console.error('DEBUG: Auth with fresh proof failed:', error);
+                        throw error;
+                    }
+                } else {
+                    console.log('DEBUG: No proof available after connection');
+                    // If we can't get proof but we're connected, try auth without it
+                    authResult = await makeAuthRequest({
+                        twa_data: WebApp.initData,
+                    });
+                }
+            } else {
+                // Case 2: Already connected - try to use stored proof first
+                console.log('DEBUG: Already connected');
+
+                // Check if we have a valid stored proof
+                const storedProofData = localStorage.getItem(tonProofStorageKey);
+                if (storedProofData) {
+                    try {
+                        const proofData = JSON.parse(storedProofData);
+
+                        // Check if the proof matches current wallet and is not too old
+                        if (tonConnectUI.wallet?.account?.address === proofData.account.address) {
+                            console.log('DEBUG: Using stored proof');
+
+                            // Try auth with stored proof but ignore errors
+                            try {
+                                authResult = await makeAuthRequest({
+                                    twa_data: WebApp.initData,
+                                    ton_proof: {
+                                        account: proofData.account,
+                                        ton_proof: proofData.proof,
+                                    },
+                                });
+
+                                // If successful, remove stored proof as it's been used
+                                localStorage.removeItem(tonProofStorageKey);
+                            } catch (error) {
+                                console.log(
+                                    'DEBUG: Auth with stored proof failed, proceeding without it'
+                                );
+
+                                // Fall back to auth without proof
+                                authResult = await makeAuthRequest({
+                                    twa_data: WebApp.initData,
+                                });
+                            }
+                        } else {
+                            console.log('DEBUG: Stored proof address mismatch');
+                            localStorage.removeItem(tonProofStorageKey);
+
+                            // Auth without proof
+                            authResult = await makeAuthRequest({
+                                twa_data: WebApp.initData,
+                            });
+                        }
+                    } catch (error) {
+                        console.error('DEBUG: Error parsing stored proof:', error);
+                        localStorage.removeItem(tonProofStorageKey);
+
+                        // Auth without proof
+                        authResult = await makeAuthRequest({
+                            twa_data: WebApp.initData,
+                        });
+                    }
+                } else {
+                    // No stored proof, check if we have a live proof
+                    if (
+                        tonConnectUI.wallet?.connectItems?.tonProof &&
+                        !('error' in tonConnectUI.wallet.connectItems.tonProof)
+                    ) {
+                        console.log('DEBUG: Using live proof from wallet');
+                        try {
+                            // Try auth with the live proof
+                            authResult = await makeAuthRequest({
+                                twa_data: WebApp.initData,
+                                ton_proof: {
+                                    account: tonConnectUI.wallet.account,
+                                    ton_proof: tonConnectUI.wallet.connectItems.tonProof.proof,
+                                },
+                            });
+                        } catch (error) {
+                            console.log(
+                                'DEBUG: Auth with live proof failed, proceeding without it'
+                            );
+
+                            // Fall back to auth without proof
+                            authResult = await makeAuthRequest({
+                                twa_data: WebApp.initData,
+                            });
+                        }
+                    } else {
+                        // Connected without proof - already authenticated
+                        console.log('DEBUG: Connected without proof, proceeding without it');
+                        authResult = await makeAuthRequest({
+                            twa_data: WebApp.initData,
+                        });
+                    }
+                }
+            }
+
+            // Always try to select wallet after auth (this validates the connection)
+            if (tonConnectUI.wallet?.account?.address) {
+                console.log('DEBUG: Selecting wallet', tonConnectUI.wallet.account.address);
+                try {
+                    await makeSelectWalletRequest({
+                        wallet_address: tonConnectUI.wallet.account.address,
+                    });
+
+                    // Additional validation check
+                    if (!isConnectionValid()) {
+                        console.log('DEBUG: Connection validation failed, disconnecting');
+                        await tonConnectUI.disconnect();
+                        localStorage.removeItem(sessionStorageKey);
+                        localStorage.removeItem(tonProofStorageKey);
+                        throw new Error('Connection validation failed');
+                    }
+                } catch (error) {
+                    // Errors from makeSelectWalletRequest are already handled
+                    console.error('DEBUG: Wallet selection failed:', error);
+                    throw error;
+                }
+            }
+
+            return authResult;
+        } catch (error) {
+            console.error('DEBUG: Auth flow failed:', error);
+            throw error;
+        } finally {
+            clearInterval(interval.current);
         }
-
-        if (tonConnectUI.wallet?.account?.address) {
-            console.log('DEBUG: Selecting wallet', tonConnectUI.wallet.account.address);
-            await makeSelectWalletRequest({ wallet_address: tonConnectUI.wallet.account.address });
-        }
-
-        return authResult;
-
-        // Commented this part for two reasons:
-        // 1) When we include ton_proof from the wallet it fails the call for a reason of bad ton_proof
-        // 2) This call could happen only if the first case happened and it means that the ton_proof is already have been stored once before
-        // Case 2: Connected with proof - use it
-        // if (wallet?.connectItems?.tonProof && !("error" in wallet.connectItems.tonProof)) {
-        //   console.log("DEBUG: Using existing proof", wallet.connectItems.tonProof.proof);
-        //   return makeAuthRequest({
-        //     twa_data: WebApp.initData,
-        //     ton_proof: {
-        //       account: wallet.account,
-        //       ton_proof: wallet.connectItems.tonProof.proof,
-        //     },
-        //   });
-        // }
     });
 };