Auto-enable jettons on restore or watch Ton blockchain

UnstoppableWallet/UnstoppableWallet/Core/Adapters/TonEventConverter.swift

@@ -38,7 +38,7 @@ class TonEventConverter {
             return .coinValue(token: token, value: value)
         } else {
             let value = convertAmount(amount: value, decimals: jetton.decimals, sign: sign)
-            return .jettonValue(name: jetton.name, symbol: jetton.symbol, decimals: jetton.decimals, value: value)
+            return .jettonValue(jetton: jetton, value: value)
         }
     }
 

UnstoppableWallet/UnstoppableWallet/Core/Adapters/TonTransactionAdapter.swift

@@ -119,7 +119,7 @@ extension TonTransactionAdapter: ITransactionsAdapter {
         tonKit.eventPublisher(tagQuery: tagQuery(token: token, filter: filter, address: address))
             .asObservable()
             .map { [converter] in
-                $0.map { converter.transactionRecord(event: $0) }
+                $0.events.map { converter.transactionRecord(event: $0) }
             }
     }
 

UnstoppableWallet/UnstoppableWallet/Core/App.swift

@@ -69,7 +69,7 @@ class App {
 
     let btcBlockchainManager: BtcBlockchainManager
     let evmSyncSourceManager: EvmSyncSourceManager
-    let evmAccountRestoreStateManager: EvmAccountRestoreStateManager
+    let restoreStateManager: RestoreStateManager
     let evmBlockchainManager: EvmBlockchainManager
     let evmLabelManager: EvmLabelManager
     let binanceKitManager: BinanceKitManager
@@ -195,10 +195,10 @@ class App {
         let evmSyncSourceStorage = EvmSyncSourceStorage(dbPool: dbPool)
         evmSyncSourceManager = EvmSyncSourceManager(testNetManager: testNetManager, blockchainSettingsStorage: blockchainSettingsStorage, evmSyncSourceStorage: evmSyncSourceStorage)
 
-        let evmAccountRestoreStateStorage = EvmAccountRestoreStateStorage(dbPool: dbPool)
-        evmAccountRestoreStateManager = EvmAccountRestoreStateManager(storage: evmAccountRestoreStateStorage)
+        let restoreStateStorage = RestoreStateStorage(dbPool: dbPool)
+        restoreStateManager = RestoreStateManager(storage: restoreStateStorage)
 
-        let evmAccountManagerFactory = EvmAccountManagerFactory(accountManager: accountManager, walletManager: walletManager, evmAccountRestoreStateManager: evmAccountRestoreStateManager, marketKit: marketKit)
+        let evmAccountManagerFactory = EvmAccountManagerFactory(accountManager: accountManager, walletManager: walletManager, restoreStateManager: restoreStateManager, marketKit: marketKit)
         evmBlockchainManager = EvmBlockchainManager(syncSourceManager: evmSyncSourceManager, testNetManager: testNetManager, marketKit: marketKit, accountManagerFactory: evmAccountManagerFactory)
 
         let hsLabelProvider = HsLabelProvider(networkManager: networkManager)
@@ -208,9 +208,9 @@ class App {
 
         binanceKitManager = BinanceKitManager()
         let tronKitManager = TronKitManager(testNetManager: testNetManager)
-        tronAccountManager = TronAccountManager(accountManager: accountManager, walletManager: walletManager, marketKit: marketKit, tronKitManager: tronKitManager, evmAccountRestoreStateManager: evmAccountRestoreStateManager)
+        tronAccountManager = TronAccountManager(accountManager: accountManager, walletManager: walletManager, marketKit: marketKit, tronKitManager: tronKitManager, restoreStateManager: restoreStateManager)
 
-        tonKitManager = TonKitManager()
+        tonKitManager = TonKitManager(restoreStateManager: restoreStateManager, marketKit: marketKit, walletManager: walletManager)
 
         let restoreSettingsStorage = RestoreSettingsStorage(dbPool: dbPool)
         restoreSettingsManager = RestoreSettingsManager(storage: restoreSettingsStorage)

UnstoppableWallet/UnstoppableWallet/Core/Factories/EvmAccountManagerFactory.swift

@@ -4,13 +4,13 @@ import MarketKit
 class EvmAccountManagerFactory {
     private let accountManager: AccountManager
     private let walletManager: WalletManager
-    private let evmAccountRestoreStateManager: EvmAccountRestoreStateManager
+    private let restoreStateManager: RestoreStateManager
     private let marketKit: MarketKit.Kit
 
-    init(accountManager: AccountManager, walletManager: WalletManager, evmAccountRestoreStateManager: EvmAccountRestoreStateManager, marketKit: MarketKit.Kit) {
+    init(accountManager: AccountManager, walletManager: WalletManager, restoreStateManager: RestoreStateManager, marketKit: MarketKit.Kit) {
         self.accountManager = accountManager
         self.walletManager = walletManager
-        self.evmAccountRestoreStateManager = evmAccountRestoreStateManager
+        self.restoreStateManager = restoreStateManager
         self.marketKit = marketKit
     }
 }
@@ -23,7 +23,7 @@ extension EvmAccountManagerFactory {
             walletManager: walletManager,
             marketKit: marketKit,
             evmKitManager: evmKitManager,
-            evmAccountRestoreStateManager: evmAccountRestoreStateManager
+            restoreStateManager: restoreStateManager
         )
     }
 }

UnstoppableWallet/UnstoppableWallet/Core/Managers/EvmAccountManager.swift

@@ -16,19 +16,19 @@ class EvmAccountManager {
     private let walletManager: WalletManager
     private let marketKit: MarketKit.Kit
     private let evmKitManager: EvmKitManager
-    private let evmAccountRestoreStateManager: EvmAccountRestoreStateManager
+    private let restoreStateManager: RestoreStateManager
 
     private let disposeBag = DisposeBag()
     private var cancellables = Set<AnyCancellable>()
     private var tasks = Set<AnyTask>()
 
-    init(blockchainType: BlockchainType, accountManager: AccountManager, walletManager: WalletManager, marketKit: MarketKit.Kit, evmKitManager: EvmKitManager, evmAccountRestoreStateManager: EvmAccountRestoreStateManager) {
+    init(blockchainType: BlockchainType, accountManager: AccountManager, walletManager: WalletManager, marketKit: MarketKit.Kit, evmKitManager: EvmKitManager, restoreStateManager: RestoreStateManager) {
         self.blockchainType = blockchainType
         self.accountManager = accountManager
         self.walletManager = walletManager
         self.marketKit = marketKit
         self.evmKitManager = evmKitManager
-        self.evmAccountRestoreStateManager = evmAccountRestoreStateManager
+        self.restoreStateManager = restoreStateManager
 
         subscribe(ConcurrentDispatchQueueScheduler(qos: .userInitiated), disposeBag, evmKitManager.evmKitCreatedObservable) { [weak self] in self?.handleEvmKitCreated() }
     }
@@ -62,7 +62,7 @@ class EvmAccountManager {
             return
         }
 
-        if initial, account.origin == .restored, !account.watchAccount, !evmAccountRestoreStateManager.isRestored(account: account, blockchainType: blockchainType) {
+        if initial, account.origin == .restored, !account.watchAccount, !restoreStateManager.shouldRestore(account: account, blockchainType: blockchainType) {
             return
         }
 

UnstoppableWallet/UnstoppableWallet/Core/Managers/RestoreStateManager.swift

@@ -0,0 +1,35 @@
+import MarketKit
+
+class RestoreStateManager {
+    private let storage: RestoreStateStorage
+
+    init(storage: RestoreStateStorage) {
+        self.storage = storage
+    }
+}
+
+extension RestoreStateManager {
+    func restoreState(account: Account, blockchainType: BlockchainType) -> RestoreState {
+        (try? storage.restoreState(accountId: account.id, blockchainUid: blockchainType.uid)) ?? RestoreState(accountId: account.id, blockchainUid: blockchainType.uid)
+    }
+
+    func shouldRestore(account: Account, blockchainType: BlockchainType) -> Bool {
+        restoreState(account: account, blockchainType: blockchainType).shouldRestore
+    }
+
+    func initialRestored(account: Account, blockchainType: BlockchainType) -> Bool {
+        restoreState(account: account, blockchainType: blockchainType).initialRestored
+    }
+
+    func setShouldRestore(account: Account, blockchainType: BlockchainType) {
+        var state = (try? storage.restoreState(accountId: account.id, blockchainUid: blockchainType.uid)) ?? RestoreState(accountId: account.id, blockchainUid: blockchainType.uid)
+        state.shouldRestore = true
+        try? storage.save(restoreState: state)
+    }
+
+    func setInitialRestored(account: Account, blockchainType: BlockchainType) {
+        var state = (try? storage.restoreState(accountId: account.id, blockchainUid: blockchainType.uid)) ?? RestoreState(accountId: account.id, blockchainUid: blockchainType.uid)
+        state.initialRestored = true
+        try? storage.save(restoreState: state)
+    }
+}

UnstoppableWallet/UnstoppableWallet/Core/Managers/TonKitManager.swift

@@ -1,21 +1,35 @@
+import Combine
 import Foundation
 import HdWalletKit
+import MarketKit
 import TonKit
 import TonSwift
 import TweetNacl
 
 class TonKitManager {
-    private weak var _tonKit: Kit?
+    private let restoreStateManager: RestoreStateManager
+    private let marketKit: MarketKit.Kit
+    private let walletManager: WalletManager
+    private var jettonBalanceCancellable: AnyCancellable?
+    private var eventCancellable: AnyCancellable?
+
+    private weak var _tonKit: TonKit.Kit?
     private var currentAccount: Account?
 
     private let queue = DispatchQueue(label: "\(AppConfig.label).ton-kit-manager", qos: .userInitiated)
 
-    private func _tonKit(account: Account) throws -> Kit {
+    init(restoreStateManager: RestoreStateManager, marketKit: MarketKit.Kit, walletManager: WalletManager) {
+        self.restoreStateManager = restoreStateManager
+        self.walletManager = walletManager
+        self.marketKit = marketKit
+    }
+
+    private func _tonKit(account: Account) throws -> TonKit.Kit {
         if let _tonKit, let currentAccount, currentAccount == account {
             return _tonKit
         }
 
-        let type: Kit.WalletType
+        let type: TonKit.Kit.WalletType
 
         switch account.type {
         case .mnemonic:
@@ -37,7 +51,7 @@ class TonKitManager {
             throw AdapterError.unsupportedAccount
         }
 
-        let tonKit = try Kit.instance(
+        let tonKit = try TonKit.Kit.instance(
             type: type,
             walletVersion: .v4,
             network: .mainNet,
@@ -51,16 +65,111 @@ class TonKitManager {
         _tonKit = tonKit
         currentAccount = account
 
+        subscribe(tonKit: tonKit, account: account)
+
         return tonKit
     }
+
+    private func subscribe(tonKit: TonKit.Kit, account: Account) {
+        let restoreState = restoreStateManager.restoreState(account: account, blockchainType: .ton)
+
+        // print("RESTORE STATE: shouldRestore: \(restoreState.shouldRestore), initialRestored: \(restoreState.initialRestored)")
+
+        if restoreState.shouldRestore || account.watchAccount, !restoreState.initialRestored {
+            jettonBalanceCancellable = tonKit.jettonBalanceMapPublisher
+                .sink { [weak self, restoreStateManager] in
+                    self?.handle(jettons: $0.values.map { $0.jetton }, account: account)
+
+                    restoreStateManager.setInitialRestored(account: account, blockchainType: .ton)
+
+                    self?.jettonBalanceCancellable?.cancel()
+                    self?.jettonBalanceCancellable = nil
+                }
+        }
+
+        let address = tonKit.receiveAddress
+
+        eventCancellable = tonKit.eventPublisher(tagQuery: .init())
+            .sink { [weak self] in self?.handle(events: $0.events, initial: $0.initial, address: address, account: account) }
+    }
+
+    private func handle(events: [Event], initial: Bool, address: TonSwift.Address, account: Account) {
+        guard !initial else {
+            // print("ignore initial events: \(events.count)")
+            return
+        }
+
+        // print("HANDLE EVENTS: \(events.count)")
+
+        var jettons = Set<Jetton>()
+
+        for event in events {
+            for action in event.actions {
+                switch action.type {
+                case let .jettonTransfer(action):
+                    if action.recipient?.address == address {
+                        jettons.insert(action.jetton)
+                    }
+                case let .jettonMint(action):
+                    if action.recipient.address == address {
+                        jettons.insert(action.jetton)
+                    }
+                case let .jettonSwap(action):
+                    if let jetton = action.jettonMasterIn {
+                        jettons.insert(jetton)
+                    }
+                default: ()
+                }
+            }
+        }
+
+        handle(jettons: Array(jettons), account: account)
+    }
+
+    private func handle(jettons: [Jetton], account: Account) {
+        // print("HANDLE JETTONS: \(jettons.map { $0.name })")
+
+        guard !jettons.isEmpty else {
+            return
+        }
+
+        let existingWallets = walletManager.activeWallets
+        let existingTokenTypeIds = existingWallets.map(\.token.type.id)
+        let newJettons = jettons.filter { !existingTokenTypeIds.contains($0.tokenType.id) }
+
+        // print("new jettons: \(newJettons.map { $0.name })")
+
+        guard !newJettons.isEmpty else {
+            return
+        }
+
+        let enabledWallets = newJettons.map { jetton in
+            EnabledWallet(
+                tokenQueryId: TokenQuery(blockchainType: .ton, tokenType: jetton.tokenType).id,
+                accountId: account.id,
+                coinName: jetton.name,
+                coinCode: jetton.symbol,
+                coinImage: jetton.image,
+                tokenDecimals: jetton.decimals
+            )
+        }
+
+        walletManager.save(enabledWallets: enabledWallets)
+    }
 }
 
 extension TonKitManager {
-    var tonKit: Kit? {
+    var tonKit: TonKit.Kit? {
         queue.sync { _tonKit }
     }
 
-    func tonKit(account: Account) throws -> Kit {
+    func tonKit(account: Account) throws -> TonKit.Kit {
         try queue.sync { try _tonKit(account: account) }
     }
 }
+
+extension Jetton {
+    var tokenType: TokenType {
+        .jetton(address: address.toString(bounceable: true))
+    }
+}

UnstoppableWallet/UnstoppableWallet/Core/Managers/TronAccountManager.swift

@@ -9,17 +9,17 @@ class TronAccountManager {
     private let walletManager: WalletManager
     private let marketKit: MarketKit.Kit
     private let tronKitManager: TronKitManager
-    private let evmAccountRestoreStateManager: EvmAccountRestoreStateManager
+    private let restoreStateManager: RestoreStateManager
 
     private let disposeBag = DisposeBag()
     private var internalDisposeBag = DisposeBag()
 
-    init(accountManager: AccountManager, walletManager: WalletManager, marketKit: MarketKit.Kit, tronKitManager: TronKitManager, evmAccountRestoreStateManager: EvmAccountRestoreStateManager) {
+    init(accountManager: AccountManager, walletManager: WalletManager, marketKit: MarketKit.Kit, tronKitManager: TronKitManager, restoreStateManager: RestoreStateManager) {
         self.accountManager = accountManager
         self.walletManager = walletManager
         self.marketKit = marketKit
         self.tronKitManager = tronKitManager
-        self.evmAccountRestoreStateManager = evmAccountRestoreStateManager
+        self.restoreStateManager = restoreStateManager
 
         subscribe(ConcurrentDispatchQueueScheduler(qos: .userInitiated), disposeBag, tronKitManager.tronKitCreatedObservable) { [weak self] in self?.handleTronKitCreated() }
     }
@@ -48,7 +48,7 @@ class TronAccountManager {
             return
         }
 
-        if initial, account.origin == .restored, !account.watchAccount, !evmAccountRestoreStateManager.isRestored(account: account, blockchainType: blockchainType) {
+        if initial, account.origin == .restored, !account.watchAccount, !restoreStateManager.shouldRestore(account: account, blockchainType: blockchainType) {
             return
         }