Merge 1789576 into d1b42a6

lib/wallet/layout.js

@@ -28,6 +28,8 @@ const bdb = require('bdb');
  *  o[hash][index] -> outpoint->wid map
  *  T[hash] -> tx->wid map
  *  t[wid]* -> txdb
+ *  N[hash256] -> name map
+ *  M[migration-id] -> dummy
  */
 
 exports.wdb = {
@@ -49,9 +51,8 @@ exports.wdb = {
   o: bdb.key('o', ['hash256', 'uint32']),
   T: bdb.key('T', ['hash256']),
   t: bdb.key('t', ['uint32']),
-
-  // Name Map
-  N: bdb.key('N', ['hash256'])
+  N: bdb.key('N', ['hash256']),
+  M: bdb.key('M', ['uint32'])
 };
 
 /*

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);

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);

lib/wallet/wallet.js

@@ -224,6 +224,48 @@ class Wallet extends EventEmitter {
     }
   }
 
+  /**
+   * Run change address migration.
+   * @param {Batch} b
+   */
+
+  async migrateChange(b) {
+    const unlock1 = await this.writeLock.lock();
+    const unlock2 = await this.fundLock.lock();
+
+    try {
+      return await this._migrateChange(b);
+    } finally {
+      unlock2();
+      unlock1();
+    }
+  }
+
+  /**
+   * Run change address migration (without a lock).
+   * @param {Batch} b
+   */
+
+  async _migrateChange(b) {
+    let total = 0;
+
+    for (let i = 0; i < this.accountDepth; i++) {
+      const account = await this.getAccount(i);
+
+      for (let i = 0; i < account.changeDepth + account.lookahead; i++) {
+        const key = account.deriveChange(i);
+        const path = key.toPath();
+
+        if (!await this.wdb.hasPath(account.wid, path.hash)) {
+          await this.wdb.savePath(b, account.wid, path);
+          total += 1;
+        }
+      }
+    }
+
+    return total;
+  }
+
   /**
    * Add a public account key to the wallet (multisig).
    * Saves the key in the wallet database.

lib/wallet/walletdb.js

@@ -206,6 +206,66 @@ class WalletDB extends EventEmitter {
       wallet.id, wallet.wid, addr.toString(this.network));
 
     this.primary = wallet;
+
+    await this.migrateChange();
+  }
+
+  /**
+   * Run change address migration.
+   * @param {Wallet} wallet
+   */
+
+  async migrateChange() {
+    if (this.options.memory)
+      return;
+
+    if (await this.db.has(layout.M.encode(0)))
+      return;
+
+    const b = this.db.batch();
+
+    b.put(layout.M.encode(0), null);
+
+    this.logger.warning('Checking change address corruption...');
+
+    const wids = await this.db.keys({
+      gte: layout.W.min(),
+      lte: layout.W.max(),
+      parse: key => layout.W.decode(key)[0]
+    });
+
+    let total = 0;
+
+    for (const wid of wids) {
+      const wallet = await this.get(wid);
+
+      this.logger.info('Checking wallet (id=%s, wid=%d).',
+                       wallet.id, wid);
+
+      total += await wallet.migrateChange(b);
+    }
+
+    if (this.options.migrate === 0 || total === 0)
+      await b.write();
+
+    if (total === 0) {
+      this.logger.info('All change addresses present.');
+      return;
+    }
+
+    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)'
+      );
+    }
   }
 
   /**
@@ -2321,6 +2381,7 @@ class WalletOptions {
 
     this.spv = false;
     this.wipeNoReally = false;
+    this.migrate = null;
 
     if (options)
       this.fromOptions(options);
@@ -2398,6 +2459,11 @@ class WalletOptions {
       this.wipeNoReally = options.wipeNoReally;
     }
 
+    if (options.migrate != null) {
+      assert((options.migrate >>> 0) === options.migrate);
+      this.migrate = options.migrate;
+    }
+
     return this;
   }
 

test/wallet-change-test.js

@@ -0,0 +1,169 @@
+/* eslint-env mocha */
+/* eslint prefer-arrow-callback: "off" */
+
+'use strict';
+
+const assert = require('bsert');
+const FullNode = require('../lib/node/fullnode');
+const Address = require('../lib/primitives/address');
+const {tmpdir} = require('os');
+const {randomBytes} = require('bcrypto/lib/random');
+const Path = require('path');
+const layouts = require('../lib/wallet/layout');
+const layout = layouts.wdb;
+
+const uniq = randomBytes(4).toString('hex');
+const path = Path.join(tmpdir(), `hsd-test-${uniq}`);
+
+const node = new FullNode({
+  prefix: path,
+  memory: false,
+  network: 'regtest',
+  plugins: [require('../lib/wallet/plugin')]
+});
+
+const {wdb} = node.require('walletdb');
+
+let wallet, recAddr;
+const changeAddrs = [];
+const manualChangeAddrs = [];
+const missingChangeAddrs = [];
+
+async function mineBlocks(n, addr) {
+  addr = addr ? addr : new Address().toString('regtest');
+  for (let i = 0; i < n; i++) {
+    const block = await node.miner.mineBlock(null, addr);
+    await node.chain.add(block);
+  }
+}
+
+describe('Derive and save change addresses', function() {
+  before(async () => {
+    await node.ensure();
+    await node.open();
+
+    wallet = await wdb.create();
+    recAddr = await wallet.receiveAddress();
+  });
+
+  after(async () => {
+    await node.close();
+  });
+
+  it('should fund account', async () => {
+    await mineBlocks(2, recAddr);
+
+    // Wallet rescan is an effective way to ensure that
+    // wallet and chain are synced before proceeding.
+    await wdb.rescan(0);
+
+    const aliceBal = await wallet.getBalance(0);
+    assert(aliceBal.confirmed === 2000 * 2 * 1e6);
+  });
+
+  it('should send 20 transactions', async () => {
+    for (let i = 0; i < 20; i++) {
+      const tx = await wallet.send({outputs: [{
+        address: Address.fromHash(Buffer.alloc(32, 1)),
+        value: 10000
+      }]});
+
+      for (const output of tx.outputs) {
+        if (output.value !== 10000)
+          changeAddrs.push(output.address);
+      }
+    }
+  });
+
+  it('should have incremented changeDepth by 20', async () => {
+    const info = await wallet.getAccount(0);
+    assert.strictEqual(info.changeDepth, 21);
+    assert.strictEqual(changeAddrs.length, 20);
+  });
+
+  it('should have all change addresses saved', async () => {
+    for (const addr of changeAddrs) {
+      assert(await wallet.hasAddress(addr));
+    }
+  });
+
+  it('should manually generate 20 change addresses', async () => {
+    for (let i = 0; i < 20; i++) {
+      const addr = await wallet.createChange();
+      manualChangeAddrs.push(addr.getAddress());
+    }
+  });
+
+  it('should have incremented changeDepth by 20', async () => {
+    const info = await wallet.getAccount(0);
+    assert.strictEqual(info.changeDepth, 41);
+    assert.strictEqual(manualChangeAddrs.length, 20);
+  });
+
+  it('should have all change addresses saved', async () => {
+    for (const addr of manualChangeAddrs) {
+      assert(await wallet.hasAddress(addr));
+    }
+  });
+
+  it('should recreate the missing change address bug', async () => {
+    for (let i = 0; i < 20; i++) {
+      const acct = await wallet.getAccount(0);
+      const key = acct.deriveChange(acct.changeDepth);
+      acct.changeDepth += 1;
+      const b = wdb.db.batch();
+      await wdb.saveAccount(b, acct);
+      await b.write();
+      missingChangeAddrs.push(key.getAddress());
+    }
+  });
+
+  it('should have no missing change addresses beyond lookahead', async () => {
+    const acct = await wallet.getAccount(0);
+    const lookahead = acct.lookahead;
+
+    for (let i = 0; i < missingChangeAddrs.length; i++) {
+      const addr = await wallet.hasAddress(missingChangeAddrs[i]);
+
+      if (i < lookahead)
+        assert(addr);
+      else
+        assert(!addr);
+    }
+  });
+
+  it('should migrate wallet and recover change addresses', async () => {
+    // Fake an old db state
+    await wdb.db.del(layout.M.encode(0));
+
+    // Run migration script without flag -- throws
+    await assert.rejects(
+      wdb.migrateChange(),
+      {
+        message: '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)'
+      }
+    );
+
+    // Add flag
+    wdb.options.migrate = 0;
+
+    // Run migration script again
+    await wdb.migrateChange();
+
+    // Fixed
+    for (const addr of missingChangeAddrs) {
+      assert(await wallet.hasAddress(addr));
+    }
+
+    // Sanity checks
+    for (const addr of changeAddrs) {
+      assert(await wallet.hasAddress(addr));
+    }
+    for (const addr of manualChangeAddrs) {
+      assert(await wallet.hasAddress(addr));
+    }
+  });
+});