From 2ff2fe3ee25b883d2ae46d9dd85d3427726287b5 Mon Sep 17 00:00:00 2001
From: Chi Kei Chan <chikeichan@gmail.com>
Date: Mon, 19 Oct 2020 18:20:37 -0700
Subject: [PATCH 1/2] wallet: remove wid from index when wallet is removed

---
 lib/wallet/walletdb.js | 10 ++++++++++
 test/wallet-test.js    | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index c5a4cc1ba..5bbc61b39 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -1157,6 +1157,16 @@ class WalletDB extends EventEmitter {
       return this.removeOutpointMap(b, hash, index, wid);
     });
 
+    const niter = this.db.iterator({
+      gte: layout.N.min(),
+      lte: layout.N.max()
+    });
+
+    await niter.each((key) => {
+      const [hash] = layout.N.decode(key);
+      return this.removeNameMap(b, hash, wid);
+    });
+
     const uiter = bucket.iterator({
       gte: tlayout.p.min(),
       lte: tlayout.p.max(),
diff --git a/test/wallet-test.js b/test/wallet-test.js
index bad336858..defad1714 100644
--- a/test/wallet-test.js
+++ b/test/wallet-test.js
@@ -29,7 +29,7 @@ const policy = require('../lib/protocol/policy');
 const HDPrivateKey = require('../lib/hd/private');
 const Wallet = require('../lib/wallet/wallet');
 const rules = require('../lib/covenants/rules');
-const {types} = rules;
+const {types, hashName} = rules;
 const {forValue} = require('./util/common');
 
 const KEY1 = 'xprv9s21ZrQH143K3Aj6xQBymM31Zb4BVc7wxqfUhMZrzewdDVCt'
@@ -1490,11 +1490,42 @@ describe('Wallet', function() {
   });
 
   it('should remove a wallet', async () => {
-    await wdb.create({
+    const wallet = await wdb.create({
       id: 'alice100'
     });
+    const addr1 = await wallet.receiveAddress();
+    const b = wdb.db.batch();
+    const wid = await wdb.getWID('alice100');
     assert(await wdb.get('alice100'));
+
+    // Add one single, unconfirmed coin to wallet
+    const mtx = new MTX();
+    mtx.addInput(dummyInput());
+    mtx.addOutput(addr1, 10 * 1e8);
+    await wdb.addTXMap(b, mtx.hash(), wid);
+
+    // Add one name to NameMap
+    await wdb.addNameMap(b, hashName('test123'), wid);
+
+    await b.write();
+
+    // Should return tx from TX Map
+    assert(await wdb.getWalletsByTX(mtx));
+
+    // Should have wid in NameMap
+    const map = await wdb.getNameMap(hashName('test123'));
+    assert(map.wids.has(wid));
+
+    // Remove wallet
     await wdb.remove('alice100');
+
+    // Should not return wid from NameMap after wid is removed
+    assert(!await wdb.getNameMap(hashName('test123')));
+
+    // Should not return tx from TX Map after wallet is removed
+    assert(!await wdb.getWalletsByTX(mtx));
+
+    // Should not return wallet after it is removed
     assert(!await wdb.get('alice100'));
   });
 

From b3c97a4c51ffb480d107bfa93ce0ebf635fb21e4 Mon Sep 17 00:00:00 2001
From: Matthew Zipkin <pinheadmz@gmail.com>
Date: Wed, 3 Feb 2021 16:14:58 -0500
Subject: [PATCH 2/2] wallet: clean up remove and test

---
 lib/wallet/walletdb.js | 20 ++++++++++----------
 test/wallet-test.js    | 15 +++++++++------
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/lib/wallet/walletdb.js b/lib/wallet/walletdb.js
index 5bbc61b39..4eba6f23d 100644
--- a/lib/wallet/walletdb.js
+++ b/lib/wallet/walletdb.js
@@ -1105,6 +1105,16 @@ class WalletDB extends EventEmitter {
       return this.removePathMap(b, hash, wid);
     });
 
+    const niter = this.db.iterator({
+      gte: layout.N.min(),
+      lte: layout.N.max()
+    });
+
+    await niter.each((key) => {
+      const [hash] = layout.N.decode(key);
+      return this.removeNameMap(b, hash, wid);
+    });
+
     const removeRange = (opt) => {
       return this.db.iterator(opt).each(key => b.del(key));
     };
@@ -1157,16 +1167,6 @@ class WalletDB extends EventEmitter {
       return this.removeOutpointMap(b, hash, index, wid);
     });
 
-    const niter = this.db.iterator({
-      gte: layout.N.min(),
-      lte: layout.N.max()
-    });
-
-    await niter.each((key) => {
-      const [hash] = layout.N.decode(key);
-      return this.removeNameMap(b, hash, wid);
-    });
-
     const uiter = bucket.iterator({
       gte: tlayout.p.min(),
       lte: tlayout.p.max(),
diff --git a/test/wallet-test.js b/test/wallet-test.js
index defad1714..57d096f70 100644
--- a/test/wallet-test.js
+++ b/test/wallet-test.js
@@ -1510,20 +1510,23 @@ describe('Wallet', function() {
     await b.write();
 
     // Should return tx from TX Map
-    assert(await wdb.getWalletsByTX(mtx));
+    let wids = await wdb.getWalletsByTX(mtx);
+    assert(wids.has(wid));
 
     // Should have wid in NameMap
-    const map = await wdb.getNameMap(hashName('test123'));
+    let map = await wdb.getNameMap(hashName('test123'));
     assert(map.wids.has(wid));
 
     // Remove wallet
     await wdb.remove('alice100');
 
-    // Should not return wid from NameMap after wid is removed
-    assert(!await wdb.getNameMap(hashName('test123')));
-
     // Should not return tx from TX Map after wallet is removed
-    assert(!await wdb.getWalletsByTX(mtx));
+    wids = await wdb.getWalletsByTX(mtx);
+    assert.strictEqual(wids, null);
+
+    // Should not return wid from NameMap after wid is removed
+    map = await wdb.getNameMap(hashName('test123'));
+    assert.strictEqual(map, null);
 
     // Should not return wallet after it is removed
     assert(!await wdb.get('alice100'));