From 11a7515bfd390608ecb05dbc0719e291e239a7c9 Mon Sep 17 00:00:00 2001
From: Christopher Jeffrey <chjjeffrey@gmail.com>
Date: Sun, 23 Oct 2016 22:24:00 -0700
Subject: [PATCH] walletdb: improve block handling.

---
 lib/wallet/walletdb.js | 51 ++++++++++++++++++------------------------
 test/wallet-test.js    | 24 +++++++++++++++-----
 2 files changed, 41 insertions(+), 34 deletions(-)

diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index e6c2a725..b3934ecf 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -1170,15 +1170,17 @@ WalletDB.prototype._rescan = co(function* rescan(chaindb, height) {
   }
 
   blocks = this.height - height;
-  assert(blocks >= 0);
+
+  if (blocks < 0)
+    throw new Error('Cannot rescan future blocks.');
 
   if (this.prune) {
-    if (blocks <= this.network.block.keepBlocks)
-      yield this.rollback(height);
-  } else {
-    yield this.rollback(height);
+    if (blocks > this.network.block.keepBlocks)
+      throw new Error('Cannot roll back beyond keepBlocks.');
   }
 
+  yield this.rollback(height);
+
   hashes = yield this.getHashes();
 
   this.logger.info('Scanning for %d addresses.', hashes.length);
@@ -1588,16 +1590,6 @@ WalletDB.prototype.rollback = co(function* rollback(height) {
 WalletDB.prototype.addBlock = co(function* addBlock(entry, txs) {
   var unlock = yield this.txLock.lock();
   try {
-    yield this.rollback(entry.height - 1);
-
-    if (entry.height <= this.height) {
-      this.logger.warning('Node is connecting low blocks in wallet.');
-      return 0;
-    }
-
-    if (entry.height !== this.height + 1)
-      throw new Error('Bad connection (height mismatch).');
-
     return yield this._addBlock(entry, txs);
   } finally {
     unlock();
@@ -1615,6 +1607,13 @@ WalletDB.prototype._addBlock = co(function* addBlock(entry, txs) {
   var total = 0;
   var i, block, tx;
 
+  // Special case for the guh-naysis block.
+  if (entry.hash === this.network.genesis.hash)
+    return;
+
+  if (entry.height !== this.height + 1)
+    throw new Error('Bad connection (height mismatch).');
+
   if (this.options.useCheckpoints) {
     if (entry.height <= this.network.checkpoints.lastHeight) {
       block = WalletBlock.fromEntry(entry);
@@ -1655,16 +1654,6 @@ WalletDB.prototype._addBlock = co(function* addBlock(entry, txs) {
 WalletDB.prototype.removeBlock = co(function* removeBlock(entry) {
   var unlock = yield this.txLock.lock();
   try {
-    yield this.rollback(entry.height);
-
-    if (entry.height > this.height) {
-      this.logger.warning('Node is disconnecting high blocks in wallet.');
-      return 0;
-    }
-
-    if (entry.height !== this.height)
-      throw new Error('Bad disconnection (height mismatch).');
-
     return yield this._removeBlock(entry);
   } finally {
     unlock();
@@ -1680,13 +1669,17 @@ WalletDB.prototype.removeBlock = co(function* removeBlock(entry) {
 
 WalletDB.prototype._removeBlock = co(function* removeBlock(entry) {
   var block = yield this.getBlock(entry.height);
-  var i, tx, prev;
+  var prev = yield this.getBlock(entry.height - 1);
+  var i, tx;
 
   if (!block)
-    return 0;
+    throw new Error('Bad disconnection (block not found).');
 
-  prev = yield this.getBlock(entry.height - 1);
-  assert(prev);
+  if (!prev)
+    throw new Error('Bad disconnection (no previous block).');
+
+  if (entry.height !== this.height)
+    throw new Error('Bad disconnection (height mismatch).');
 
   if (block.txs.length === 0) {
     yield this.setBlock(prev);
diff --git a/test/wallet-test.js b/test/wallet-test.js
index 8a2f216c..65704b06 100644
--- a/test/wallet-test.js
+++ b/test/wallet-test.js
@@ -19,12 +19,26 @@ var KEY2 = 'xprv9s21ZrQH143K3mqiSThzPtWAabQ22Pjp3uSNnZ53A5bQ4udp'
   + 'faKekc2m4AChLYH1XDzANhrSdxHYWUeTWjYJwFwWFyHkTMnMeAcW4JyRCZa';
 
 var globalHeight = 1;
+var globalTime = utils.now();
 
 function nextBlock(height) {
-  var hash;
-  height = height != null ? height : globalHeight++;
+  var hash, prev;
+
+  if (height == null)
+    height = globalHeight++;
+
   hash = crypto.hash256(utils.U32(height)).toString('hex');
-  return new WalletBlock(hash, height++, utils.now());
+  prev = crypto.hash256(utils.U32(height - 1)).toString('hex');
+
+  return {
+    hash: hash,
+    height: height,
+    prevBlock: prev,
+    ts: globalTime + height,
+    merkleRoot: constants.NULL_HASH,
+    nonce: 0,
+    bits: 0
+  };
 }
 
 function dummy(hash) {
@@ -709,9 +723,9 @@ describe('Wallet', function() {
 
     // Simulate a confirmation
     var block = nextBlock();
-    utx.ts = block.ts;
     utx.height = block.height;
     utx.block = block.hash;
+    utx.ts = block.ts;
     utx.index = 0;
 
     assert.equal(w1[depth], 1);
@@ -752,9 +766,9 @@ describe('Wallet', function() {
 
     // Simulate a confirmation
     var block = nextBlock();
-    send.ts = block.ts;
     send.height = block.height;
     send.block = block.hash;
+    send.ts = block.ts;
     send.index = 0;
 
     yield walletdb.addBlock(block, [send]);