From 223c4811804e495b2ea005d060682aa70f826261 Mon Sep 17 00:00:00 2001
From: Matthew Zipkin <pinheadmz@gmail.com>
Date: Fri, 7 Aug 2020 15:00:44 -0400
Subject: [PATCH] wallet: migrate change and rescan using config option

---
 lib/wallet/node.js     |  3 ++-
 lib/wallet/plugin.js   |  3 ++-
 lib/wallet/walletdb.js | 35 +++++++++++++++++++++++------------
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/lib/wallet/node.js b/lib/wallet/node.js
index 8c8450dc6..a9725174e 100644
--- a/lib/wallet/node.js
+++ b/lib/wallet/node.js
@@ -50,7 +50,8 @@ class WalletNode extends Node {
       maxFiles: this.config.uint('max-files'),
       cacheSize: this.config.mb('cache-size'),
       wipeNoReally: this.config.bool('wipe-no-really'),
-      spv: this.config.bool('spv')
+      spv: this.config.bool('spv'),
+      migrate: this.config.uint('migrate')
     });
 
     this.rpc = new RPC(this);
diff --git a/lib/wallet/plugin.js b/lib/wallet/plugin.js
index c0bdb585e..59ce7d184 100644
--- a/lib/wallet/plugin.js
+++ b/lib/wallet/plugin.js
@@ -51,7 +51,8 @@ class Plugin extends EventEmitter {
       maxFiles: this.config.uint('max-files'),
       cacheSize: this.config.mb('cache-size'),
       wipeNoReally: this.config.bool('wipe-no-really'),
-      spv: node.spv
+      spv: node.spv,
+      migrate: this.config.uint('migrate')
     });
 
     this.rpc = new RPC(this);
diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index 93c84efd8..3c9b5b769 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -226,12 +226,7 @@ class WalletDB extends EventEmitter {
 
     b.put(layout.M.encode(0), null);
 
-    if (!this.state.marked) {
-      await b.write();
-      return;
-    }
-
-    this.logger.warning('Running change address migration.');
+    this.logger.warning('Checking change address corruption...');
 
     const wids = await this.db.keys({
       gte: layout.W.min(),
@@ -244,19 +239,29 @@ class WalletDB extends EventEmitter {
     for (const wid of wids) {
       const wallet = await this.get(wid);
 
-      this.logger.info('Regenerating change addresses (id=%s, wid=%d).',
+      this.logger.info('Checking wallet (id=%s, wid=%d).',
                        wallet.id, wid);
 
       total += await wallet.migrateChange(b);
     }
 
-    await b.write();
+    if (this.options.migrate === 0)
+      await b.write();
 
     if (total > 0) {
-      this.logger.warning('Regenerated %d change addresses.', total);
-      this.logger.warning('Rescanning...');
-
-      await this.scan();
+      if (this.options.migrate === 0) {
+        this.logger.warning('Regenerated %d change addresses.', total);
+        this.logger.warning('Rescanning...');
+
+        await this.scan();
+      } else {
+        throw new Error(
+          'Wallet is corrupted.\n' +
+          'Back up wallet and then restart with\n' +
+          '`hsd --wallet-migrate=0` or `hs-wallet --migrate=0`\n' +
+          '(Full node required)'
+        );
+      }
     } else {
       this.logger.info('All change addresses present.');
     }
@@ -2375,6 +2380,7 @@ class WalletOptions {
 
     this.spv = false;
     this.wipeNoReally = false;
+    this.migrate = null;
 
     if (options)
       this.fromOptions(options);
@@ -2452,6 +2458,11 @@ class WalletOptions {
       this.wipeNoReally = options.wipeNoReally;
     }
 
+    if (options.migrate != null) {
+      assert((options.migrate >>> 0) === options.migrate);
+      this.migrate = options.migrate;
+    }
+
     return this;
   }