From b5bc5ff9ed6252c7e8b3203af04d99d90f1dbb22 Mon Sep 17 00:00:00 2001 From: ThomasV Date: Fri, 18 Aug 2023 15:13:33 +0200 Subject: [PATCH] Separate WalletDB from storage upgrades. Make sure that WalletDB.data is always a StoredDict. Perform db upgrades in a separate class, since they operate on a dict object. --- electrum/daemon.py | 9 +- electrum/gui/qml/qewalletdb.py | 17 ++- electrum/gui/qt/__init__.py | 9 +- electrum/gui/qt/wizard/wallet.py | 30 +++-- electrum/json_db.py | 30 ++--- electrum/lnwatcher.py | 2 +- electrum/tests/test_storage_upgrade.py | 43 ++++--- electrum/tests/test_wallet.py | 8 +- electrum/wallet.py | 4 +- electrum/wallet_db.py | 155 ++++++++++++++----------- 10 files changed, 160 insertions(+), 147 deletions(-) diff --git a/electrum/daemon.py b/electrum/daemon.py index 553b2a0ef..b21ccc128 100644 --- a/electrum/daemon.py +++ b/electrum/daemon.py @@ -47,7 +47,7 @@ from .util import log_exceptions, ignore_exceptions, randrange, OldTaskGroup from .util import EventListener, event_listener from .wallet import Wallet, Abstract_Wallet from .storage import WalletStorage -from .wallet_db import WalletDB +from .wallet_db import WalletDB, WalletRequiresSplit, WalletRequiresUpgrade from .commands import known_commands, Commands from .simple_config import SimpleConfig from .exchange_rate import FxThread @@ -512,10 +512,11 @@ class Daemon(Logger): return storage.decrypt(password) # read data, pass it to db - db = WalletDB(storage.read(), storage=storage, manual_upgrades=manual_upgrades) - if db.requires_split(): + try: + db = WalletDB(storage.read(), storage=storage, manual_upgrades=manual_upgrades) + except WalletRequiresSplit: return - if db.requires_upgrade(): + except WalletRequiresUpgrade: return if db.get_action(): return diff --git a/electrum/gui/qml/qewalletdb.py b/electrum/gui/qml/qewalletdb.py index bbf91c3e3..ea5424251 100644 --- a/electrum/gui/qml/qewalletdb.py +++ b/electrum/gui/qml/qewalletdb.py @@ -5,7 +5,7 @@ from PyQt5.QtCore import pyqtProperty, pyqtSignal, pyqtSlot, QObject from electrum.i18n import _ from electrum.logging import get_logger from electrum.storage import WalletStorage -from electrum.wallet_db import WalletDB +from electrum.wallet_db import WalletDB, WalletRequiresSplit from electrum.wallet import Wallet from electrum.util import InvalidPassword, WalletFileException, send_exception_to_crash_reporter @@ -143,7 +143,10 @@ class QEWalletDB(QObject): else: # storage not encrypted; but it might still have a keystore pw # FIXME hack... load both db and full wallet, just to tell if it has keystore pw. # this also completely ignores db.requires_split(), db.get_action(), etc - db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False) + try: + db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False) + except WalletRequiresSplit as e: + raise WalletFileException(_('This wallet requires to be split. This is currently not supported on mobile')) wallet = Wallet(db, config=self._config) self.needsPassword = wallet.has_password() if self.needsPassword: @@ -162,18 +165,14 @@ class QEWalletDB(QObject): def _load_db(self): """can raise WalletFileException""" # needs storage accessible - self._db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=True) - if self._db.requires_split(): + try: + self._db = WalletDB(self._storage.read(), storage=self._storage, manual_upgrades=False) + except WalletRequiresSplit as e: self._logger.warning('wallet requires split') raise WalletFileException(_('This wallet needs splitting. This is not supported on mobile')) if self._db.get_action(): self._logger.warning('action pending. QML version doesn\'t support continuation of wizard') raise WalletFileException(_('This wallet has an action pending. This is currently not supported on mobile')) - if self._db.requires_upgrade(): - self._logger.warning('wallet requires upgrade, upgrading') - self._db.upgrade() - self._db.write() - self._ready = True self.readyChanged.emit() diff --git a/electrum/gui/qt/__init__.py b/electrum/gui/qt/__init__.py index 376a4a06f..f00073cf2 100644 --- a/electrum/gui/qt/__init__.py +++ b/electrum/gui/qt/__init__.py @@ -61,7 +61,7 @@ from electrum.plugin import run_hook from electrum.util import (UserCancelled, profiler, send_exception_to_crash_reporter, WalletFileException, BitcoinException, get_new_wallet_name) from electrum.wallet import Wallet, Abstract_Wallet -from electrum.wallet_db import WalletDB +from electrum.wallet_db import WalletDB, WalletRequiresSplit, WalletRequiresUpgrade from electrum.logging import Logger from electrum.gui import BaseElectrumGui from electrum.simple_config import SimpleConfig @@ -430,10 +430,11 @@ class ElectrumGui(BaseElectrumGui, Logger): if storage.is_encrypted_with_user_pw() or storage.is_encrypted_with_hw_device(): storage.decrypt(d['password']) - db = WalletDB(storage.read(), storage=storage, manual_upgrades=True) - if db.requires_split() or db.requires_upgrade(): + try: + db = WalletDB(storage.read(), storage=storage, manual_upgrades=True) + except WalletRequiresSplit as e: try: - wizard.run_upgrades(db) + wizard.run_split(storage, e._split_data) except UserCancelled: return diff --git a/electrum/gui/qt/wizard/wallet.py b/electrum/gui/qt/wizard/wallet.py index 931c7d107..d2ecf8800 100644 --- a/electrum/gui/qt/wizard/wallet.py +++ b/electrum/gui/qt/wizard/wallet.py @@ -2,6 +2,7 @@ import os import sys import threading import time +import json from typing import TYPE_CHECKING @@ -21,6 +22,7 @@ from electrum.wallet import wallet_types from .wizard import QEAbstractWizard, WizardComponent from electrum.logging import get_logger, Logger from electrum import WalletStorage, mnemonic, keystore +from electrum.wallet_db import WalletDB from electrum.wizard import NewWalletWizard from electrum.gui.qt.bip39_recovery_dialog import Bip39RecoveryDialog @@ -34,7 +36,6 @@ if TYPE_CHECKING: from electrum.simple_config import SimpleConfig from electrum.plugin import Plugins from electrum.daemon import Daemon - from electrum.wallet_db import WalletDB from electrum.gui.qt import QElectrumApplication WIF_HELP_TEXT = (_('WIF keys are typed in Electrum, based on script type.') + '\n\n' + @@ -161,25 +162,20 @@ class QENewWalletWizard(NewWalletWizard, QEAbstractWizard, MessageBoxMixin): self._password = data['password'] self.path = path - def run_upgrades(self, db: 'WalletDB') -> None: - storage = db.storage - path = storage.path - if db.requires_split(): - msg = _( - "The wallet '{}' contains multiple accounts, which are no longer supported since Electrum 2.7.\n\n" - "Do you want to split your wallet into multiple files?").format(path) - if not self.question(msg): - return - file_list = db.split_accounts(path) + def run_split(self, storage, split_data) -> None: + root_path = storage.path + msg = _( + "The wallet '{}' contains multiple accounts, which are no longer supported since Electrum 2.7.\n\n" + "Do you want to split your wallet into multiple files?").format(root_path) + if self.question(msg): + file_list = WalletDB.split_accounts(root_path, split_data) msg = _('Your accounts have been moved to') + ':\n' + '\n'.join(file_list) + '\n\n' + _( - 'Do you want to delete the old file') + ':\n' + path + 'Do you want to delete the old file') + ':\n' + root_path if self.question(msg): - os.remove(path) + os.remove(root_path) self.show_warning(_('The file was removed')) - # raise now, to avoid having the old storage opened - raise UserCancelled() - if db.requires_upgrade(): - self.waiting_dialog(db.upgrade, _('Upgrading wallet format...')) + # raise now, to avoid having the old storage opened + raise UserCancelled() def is_finalized(self, wizard_data: dict) -> bool: # check decryption of existing wallet and keep wizard open if incorrect. diff --git a/electrum/json_db.py b/electrum/json_db.py index 91911ca50..341ba642d 100644 --- a/electrum/json_db.py +++ b/electrum/json_db.py @@ -34,8 +34,6 @@ from .logging import Logger if TYPE_CHECKING: from .storage import WalletStorage -JsonDBJsonEncoder = util.MyEncoder - def modifier(func): def wrapper(self, *args, **kwargs): with self.lock: @@ -168,24 +166,28 @@ class StoredDict(dict): class JsonDB(Logger): - def __init__(self, data, storage=None): + def __init__(self, s: str, storage=None, encoder=None): Logger.__init__(self) self.lock = threading.RLock() self.storage = storage + self.encoder = encoder self._modified = False # load data - if data: - self.load_data(data) - else: - self.data = {} - - def load_data(self, s): + data = self.load_data(s) + # convert to StoredDict + self.data = StoredDict(data, self, []) + + def load_data(self, s:str) -> dict: + """ overloaded in wallet_db """ + if s == '': + return {} try: - self.data = json.loads(s) + data = json.loads(s) except Exception: raise WalletFileException("Cannot read wallet file. (parsing failed)") - if not isinstance(self.data, dict): + if not isinstance(data, dict): raise WalletFileException("Malformed wallet file (not dict)") + return data def set_modified(self, b): with self.lock: @@ -204,8 +206,8 @@ class JsonDB(Logger): @modifier def put(self, key, value): try: - json.dumps(key, cls=JsonDBJsonEncoder) - json.dumps(value, cls=JsonDBJsonEncoder) + json.dumps(key, cls=self.encoder) + json.dumps(value, cls=self.encoder) except Exception: self.logger.info(f"json error: cannot save {repr(key)} ({repr(value)})") return False @@ -235,7 +237,7 @@ class JsonDB(Logger): self.data, indent=4 if human_readable else None, sort_keys=bool(human_readable), - cls=JsonDBJsonEncoder, + cls=self.encoder, ) def _should_convert_to_stored_dict(self, key) -> bool: diff --git a/electrum/lnwatcher.py b/electrum/lnwatcher.py index 627fe49db..d2870998e 100644 --- a/electrum/lnwatcher.py +++ b/electrum/lnwatcher.py @@ -325,7 +325,7 @@ class WatchTower(LNWatcher): LOGGING_SHORTCUT = 'W' def __init__(self, network: 'Network'): - adb = AddressSynchronizer(WalletDB({}, storage=None, manual_upgrades=False), network.config, name=self.diagnostic_name()) + adb = AddressSynchronizer(WalletDB('', storage=None, manual_upgrades=False), network.config, name=self.diagnostic_name()) adb.start_network(network) LNWatcher.__init__(self, adb, network) self.network = network diff --git a/electrum/tests/test_storage_upgrade.py b/electrum/tests/test_storage_upgrade.py index b12649106..85ec86787 100644 --- a/electrum/tests/test_storage_upgrade.py +++ b/electrum/tests/test_storage_upgrade.py @@ -7,7 +7,7 @@ import asyncio import inspect import electrum -from electrum.wallet_db import WalletDB +from electrum.wallet_db import WalletDBUpgrader, WalletDB, WalletRequiresUpgrade, WalletRequiresSplit from electrum.wallet import Wallet from electrum import constants from electrum import util @@ -331,31 +331,30 @@ class TestStorageUpgrade(WalletTestCase): async def _upgrade_storage(self, wallet_json, accounts=1) -> Optional[WalletDB]: if accounts == 1: # test manual upgrades - db = self._load_db_from_json_string(wallet_json=wallet_json, - manual_upgrades=True) - self.assertFalse(db.requires_split()) - if db.requires_upgrade(): - db.upgrade() + try: + db = self._load_db_from_json_string( + wallet_json=wallet_json, + manual_upgrades=True) + except WalletRequiresUpgrade: + db = self._load_db_from_json_string( + wallet_json=wallet_json, + manual_upgrades=False) await self._sanity_check_upgraded_db(db) - # test automatic upgrades - db2 = self._load_db_from_json_string(wallet_json=wallet_json, - manual_upgrades=False) - await self._sanity_check_upgraded_db(db2) - return db2 + return db else: - db = self._load_db_from_json_string(wallet_json=wallet_json, - manual_upgrades=True) - self.assertTrue(db.requires_split()) - split_data = db.get_split_accounts() - self.assertEqual(accounts, len(split_data)) - for item in split_data: - data = json.dumps(item) - new_db = WalletDB(data, storage=None, manual_upgrades=False) - await self._sanity_check_upgraded_db(new_db) + try: + db = self._load_db_from_json_string( + wallet_json=wallet_json, + manual_upgrades=True) + except WalletRequiresSplit as e: + split_data = e._split_data + self.assertEqual(accounts, len(split_data)) + for item in split_data: + data = json.dumps(item) + new_db = WalletDB(data, storage=None, manual_upgrades=False) + await self._sanity_check_upgraded_db(new_db) async def _sanity_check_upgraded_db(self, db): - self.assertFalse(db.requires_split()) - self.assertFalse(db.requires_upgrade()) wallet = Wallet(db, config=self.config) await wallet.stop() diff --git a/electrum/tests/test_wallet.py b/electrum/tests/test_wallet.py index 6bc9bef00..a1bd9b7d7 100644 --- a/electrum/tests/test_wallet.py +++ b/electrum/tests/test_wallet.py @@ -15,7 +15,7 @@ from electrum.wallet import (Abstract_Wallet, Standard_Wallet, create_new_wallet from electrum.exchange_rate import ExchangeBase, FxThread from electrum.util import TxMinedInfo, InvalidPassword from electrum.bitcoin import COIN -from electrum.wallet_db import WalletDB +from electrum.wallet_db import WalletDB, JsonDB from electrum.simple_config import SimpleConfig from electrum import util @@ -60,14 +60,14 @@ class TestWalletStorage(WalletTestCase): contents = f.write(contents) storage = WalletStorage(self.wallet_path) - db = WalletDB(storage.read(), storage=storage, manual_upgrades=True) + db = JsonDB(storage.read(), storage=storage) self.assertEqual("b", db.get("a")) self.assertEqual("d", db.get("c")) def test_write_dictionary_to_file(self): storage = WalletStorage(self.wallet_path) - db = WalletDB('', storage=storage, manual_upgrades=True) + db = JsonDB('', storage=storage) some_dict = { u"a": u"b", @@ -110,7 +110,7 @@ class FakeWallet: def __init__(self, fiat_value): super().__init__() self.fiat_value = fiat_value - self.db = WalletDB("{}", storage=None, manual_upgrades=True) + self.db = WalletDB('', storage=None, manual_upgrades=True) self.adb = FakeADB() self.db.transactions = self.db.verified_tx = {'abc':'Tx'} diff --git a/electrum/wallet.py b/electrum/wallet.py index cfef599f8..bb38cbe53 100644 --- a/electrum/wallet.py +++ b/electrum/wallet.py @@ -309,8 +309,8 @@ class Abstract_Wallet(ABC, Logger, EventListener): def __init__(self, db: WalletDB, *, config: SimpleConfig): - if not db.is_ready_to_be_used_by_wallet(): - raise Exception("storage not ready to be used by Abstract_Wallet") + #if not db.is_ready_to_be_used_by_wallet(): + # raise Exception("storage not ready to be used by Abstract_Wallet") self.config = config assert self.config is not None, "config must not be None" diff --git a/electrum/wallet_db.py b/electrum/wallet_db.py index 01a4c6f49..2d26df963 100644 --- a/electrum/wallet_db.py +++ b/electrum/wallet_db.py @@ -36,7 +36,7 @@ import time import attr from . import util, bitcoin -from .util import profiler, WalletFileException, multisig_type, TxMinedInfo, bfh +from .util import profiler, WalletFileException, multisig_type, TxMinedInfo, bfh, MyEncoder from .invoices import Invoice, Request from .keystore import bip44_derivation from .transaction import Transaction, TxOutpoint, tx_from_any, PartialTransaction, PartialTxOutput @@ -49,6 +49,11 @@ from .plugin import run_hook, plugin_loaders from .version import ELECTRUM_VERSION +class WalletRequiresUpgrade(WalletFileException): + pass +class WalletRequiresSplit(WalletFileException): + def __init__(self, split_data): + self._split_data = split_data # seed_version is now used for the version of the wallet file @@ -100,46 +105,21 @@ for key in ['locked_in', 'fails', 'settles']: json_db.register_parent_key(key, lambda x: HTLCOwner(int(x))) -class WalletDB(JsonDB): - def __init__(self, data, *, storage=None, manual_upgrades: bool): - JsonDB.__init__(self, data, storage) - if not data: - # create new DB - self.put('seed_version', FINAL_SEED_VERSION) - self._add_db_creation_metadata() - self._after_upgrade_tasks() - self._manual_upgrades = manual_upgrades - self._called_after_upgrade_tasks = False - if not self._manual_upgrades and self.requires_split(): - raise WalletFileException("This wallet has multiple accounts and must be split") - if not self.requires_upgrade(): - self._after_upgrade_tasks() - elif not self._manual_upgrades: - self.upgrade() - # load plugins that are conditional on wallet type - self.load_plugins() +class WalletDBUpgrader(Logger): - def load_data(self, s): - try: - JsonDB.load_data(self, s) - except Exception: - try: - d = ast.literal_eval(s) - labels = d.get('labels', {}) - except Exception as e: - raise WalletFileException("Cannot read wallet file. (parsing failed)") - self.data = {} - for key, value in d.items(): - try: - json.dumps(key) - json.dumps(value) - except Exception: - self.logger.info(f'Failed to convert label to json format: {key}') - continue - self.data[key] = value - if not isinstance(self.data, dict): - raise WalletFileException("Malformed wallet file (not dict)") + def __init__(self, data): + Logger.__init__(self) + self.data = data + + def get(self, key, default=None): + return self.data.get(key, default) + + def put(self, key, value): + if value is not None: + self.data[key] = value + else: + self.data.pop(key, None) def requires_split(self): d = self.get('accounts', {}) @@ -192,9 +172,6 @@ class WalletDB(JsonDB): @profiler def upgrade(self): self.logger.info('upgrading wallet format') - if self._called_after_upgrade_tasks: - # we need strict ordering between upgrade() and after_upgrade_tasks() - raise Exception("'after_upgrade_tasks' must NOT be called before 'upgrade'") self._convert_imported() self._convert_wallet_type() self._convert_account() @@ -242,12 +219,6 @@ class WalletDB(JsonDB): self._convert_version_54() self.put('seed_version', FINAL_SEED_VERSION) # just to be sure - self._after_upgrade_tasks() - - def _after_upgrade_tasks(self): - self._called_after_upgrade_tasks = True - self._load_transactions() - def _convert_wallet_type(self): if not self._is_upgrade_method_needed(0, 13): return @@ -1140,7 +1111,6 @@ class WalletDB(JsonDB): else: return True - @locked def get_seed_version(self): seed_version = self.get('seed_version') if not seed_version: @@ -1191,14 +1161,62 @@ class WalletDB(JsonDB): # generic exception raise WalletFileException(msg) - def _add_db_creation_metadata(self): - # store this for debugging purposes - v = DBMetadata( - creation_timestamp=int(time.time()), - first_electrum_version_used=ELECTRUM_VERSION, - ) - assert self.get("db_metadata", None) is None - self.put("db_metadata", v) + +class WalletDB(JsonDB): + + def __init__(self, s, *, storage=None, manual_upgrades=True): + self._upgrade = not manual_upgrades + JsonDB.__init__(self, s, storage, encoder=MyEncoder) + # create pointers + self.load_transactions() + # load plugins that are conditional on wallet type + self.load_plugins() + + def load_data(self, s): + try: + data = JsonDB.load_data(self, s) + except Exception: + try: + d = ast.literal_eval(s) + labels = d.get('labels', {}) + except Exception as e: + raise WalletFileException("Cannot read wallet file. (parsing failed)") + data = {} + for key, value in d.items(): + try: + json.dumps(key) + json.dumps(value) + except Exception: + self.logger.info(f'Failed to convert label to json format: {key}') + continue + data[key] = value + if not isinstance(data, dict): + raise WalletFileException("Malformed wallet file (not dict)") + + if len(data) == 0: + # create new DB + data['seed_version'] = FINAL_SEED_VERSION + # store this for debugging purposes + v = DBMetadata( + creation_timestamp=int(time.time()), + first_electrum_version_used=ELECTRUM_VERSION, + ) + assert data.get("db_metadata", None) is None + data["db_metadata"] = v + + dbu = WalletDBUpgrader(data) + if dbu.requires_split(): + raise WalletRequiresSplit(dbu.get_split_accounts()) + if self._upgrade: + dbu.upgrade() + if dbu.requires_upgrade(): + raise WalletRequiresUpgrade() + return dbu.data + + + @locked + def get_seed_version(self): + return self.get('seed_version') def get_db_metadata(self) -> Optional[DBMetadata]: # field only present for wallet files created with ver 4.4.0 or later @@ -1560,8 +1578,7 @@ class WalletDB(JsonDB): self._addr_to_addr_index[addr] = (1, i) @profiler - def _load_transactions(self): - self.data = StoredDict(self.data, self, []) + def load_transactions(self): # references in self.data # TODO make all these private # txid -> address -> prev_outpoint -> value @@ -1608,21 +1625,19 @@ class WalletDB(JsonDB): return True def is_ready_to_be_used_by_wallet(self): - return not self.requires_upgrade() and self._called_after_upgrade_tasks + return not self._requires_upgrade - def split_accounts(self, root_path): + @classmethod + def split_accounts(klass, root_path, split_data): from .storage import WalletStorage - out = [] - result = self.get_split_accounts() - for data in result: + file_list = [] + for data in split_data: path = root_path + '.' + data['suffix'] - storage = WalletStorage(path) - db = WalletDB(json.dumps(data), storage=storage, manual_upgrades=False) - db._called_after_upgrade_tasks = False - db.upgrade() + item_storage = WalletStorage(path) + db = WalletDB(json.dumps(data), storage=item_storage, manual_upgrades=False) db.write() - out.append(path) - return out + file_list.append(path) + return file_list def get_action(self): action = run_hook('get_action', self)