wallet: improve lookahead logic & make rpc persistent

5 files changed, 147 insertions, 31 deletions

diff --git a/src/wallet/wallet2.cpp b/src/wallet/wallet2.cpp
index 808a6fdc2..64dfc2d3b 100644
--- a/src/wallet/wallet2.cpp
+++ b/src/wallet/wallet2.cpp
@@ -1965,26 +1965,41 @@ void wallet2::set_subaddress_lookahead(size_t major, size_t minor)
   THROW_WALLET_EXCEPTION_IF(minor == 0, error::wallet_internal_error, "Subaddress minor lookahead may not be zero");
   THROW_WALLET_EXCEPTION_IF(minor > 0xffffffff, error::wallet_internal_error, "Subaddress minor lookahead is too large");
 
-  if (major > m_subaddress_lookahead_major) { // if increasing the lookahead
-    // then generate new subaddress pubkeys and add them to m_subaddresses table
-    for (uint32_t i = m_subaddress_labels.size()+m_subaddress_lookahead_major; i < m_subaddress_labels.size()+major; i++) { // m_subaddress_labels are the accounts the user is conciously keeping track of. We want that number plus the lookahead major accounts in our key table
-      for (uint32_t j = 0; j < minor; j++) { // these are newly made accounts, minor index will start from zero
-        cryptonote::subaddress_index idx = {i,j};
-        create_one_off_subaddress(idx); // then generate the key and add it to the table
-      }
-    }
-  }
-  if (minor > m_subaddress_lookahead_minor) { // if increasing the minor lookahead we need to also go back and expand the existing accounts
-    for (uint32_t i = 0; i < m_subaddress_labels.size()+m_subaddress_lookahead_major; i++) {
-      uint32_t minor_idx_start = i < m_subaddress_labels.size() ? m_subaddress_labels[i].size()+m_subaddress_lookahead_minor : m_subaddress_lookahead_minor; // if there are existing minor indices being tracked under this account we need to account for that
-      for (uint32_t j = minor_idx_start; j < minor; j++) {
-        cryptonote::subaddress_index idx = {i,j};
-        create_one_off_subaddress(idx);
-      }
-    }
-  }
+  const uint32_t old_major_lookahead = m_subaddress_lookahead_major;
+  const uint32_t old_minor_lookahead = m_subaddress_lookahead_minor;
+
   m_subaddress_lookahead_major = major;
   m_subaddress_lookahead_minor = minor;
+
+  if (old_major_lookahead >= major && old_minor_lookahead >= minor)
+    return;
+
+  // Expand the subaddresses map so that outputs received to the higher lookaheads will be identified in the scan loop
+  hw::device &hwdev = m_account.get_device();
+  cryptonote::subaddress_index index2;
+  const uint32_t max_major_idx = this->get_num_subaddress_accounts() > 0 ? (this->get_num_subaddress_accounts() - 1) : 0;
+  const uint32_t major_end = get_subaddress_clamped_sum(max_major_idx, major);
+  for (index2.major = 0; index2.major < major_end; ++index2.major)
+  {
+    // The existing minor addresses already set for this account
+    const uint32_t n_minor_subaddrs = this->get_num_subaddresses(index2.major);
+
+    // The subaddress lookahead is expected to expand from the max index in expand_subaddresses
+    const uint32_t max_minor_idx = n_minor_subaddrs > 0 ? (n_minor_subaddrs - 1) : 0;
+    const uint32_t begin = (n_minor_subaddrs || index2.major < old_major_lookahead) ? get_subaddress_clamped_sum(max_minor_idx, old_minor_lookahead) : 0;
+    // The expected new n minor subaddresses allocated for this account
+    const uint32_t end = get_subaddress_clamped_sum(max_minor_idx, minor);
+
+    if (begin >= end)
+      continue;
+
+    const std::vector<crypto::public_key> pkeys = hwdev.get_subaddress_spend_public_keys(m_account.get_keys(), index2.major, begin, end);
+    for (index2.minor = begin; index2.minor < end; ++index2.minor)
+    {
+      const crypto::public_key &D = pkeys.at(index2.minor - begin);
+      m_subaddresses[D] = index2;
+    }
+  }
 }
 //----------------------------------------------------------------------------------------------------
 /*!
diff --git a/src/wallet/wallet_rpc_server.cpp b/src/wallet/wallet_rpc_server.cpp
index 37e0968b8..15ca1500e 100644
--- a/src/wallet/wallet_rpc_server.cpp
+++ b/src/wallet/wallet_rpc_server.cpp
@@ -653,11 +653,24 @@ namespace tools
   bool wallet_rpc_server::on_set_subaddr_lookahead(const wallet_rpc::COMMAND_RPC_SET_SUBADDR_LOOKAHEAD::request& req, wallet_rpc::COMMAND_RPC_SET_SUBADDR_LOOKAHEAD::response& res, epee::json_rpc::error& er, const connection_context *ctx)
   {
     if (!m_wallet) return not_open(er);
-    try {
-      m_wallet->set_subaddress_lookahead(req.major_idx, req.minor_idx);
+    CHECK_IF_BACKGROUND_SYNCING();
+    const std::string wallet_file = m_wallet->get_wallet_file();
+    if (wallet_file == "" || m_wallet->verify_password(req.password))
+    {
+      try
+      {
+        m_wallet->set_subaddress_lookahead(req.major_idx, req.minor_idx);
+        m_wallet->rewrite(wallet_file, req.password);
+      }
+      catch (const std::exception& e) {
+        handle_rpc_exception(std::current_exception(), er, WALLET_RPC_ERROR_CODE_UNKNOWN_ERROR);
+        return false;
+      }
     }
-    catch (const std::exception& e) {
-      handle_rpc_exception(std::current_exception(), er, WALLET_RPC_ERROR_CODE_UNKNOWN_ERROR);
+    else
+    {
+      er.code = WALLET_RPC_ERROR_CODE_INVALID_PASSWORD;
+      er.message = "Invalid password.";
       return false;
     }
     return true;
diff --git a/src/wallet/wallet_rpc_server_commands_defs.h b/src/wallet/wallet_rpc_server_commands_defs.h
index 4a11845d4..ebcd1135c 100644
--- a/src/wallet/wallet_rpc_server_commands_defs.h
+++ b/src/wallet/wallet_rpc_server_commands_defs.h
@@ -186,9 +186,11 @@ namespace wallet_rpc
   {
     struct request_t
     {
+      std::string password;
       uint32_t major_idx;
       uint32_t minor_idx;
       BEGIN_KV_SERIALIZE_MAP()
+        KV_SERIALIZE(password)
         KV_SERIALIZE(major_idx)
         KV_SERIALIZE(minor_idx)
       END_KV_SERIALIZE_MAP()
diff --git a/tests/unit_tests/subaddress.cpp b/tests/unit_tests/subaddress.cpp
index b03e58f5a..71cbdbee6 100644
--- a/tests/unit_tests/subaddress.cpp
+++ b/tests/unit_tests/subaddress.cpp
@@ -103,12 +103,36 @@ TEST_F(WalletSubaddress, OutOfBoundsIndexes)
   }   
 }
 
-TEST_F(WalletSubaddress, ExpandPubkeyTable)
+// Helper function to check max subaddrs allocated
+static void check_expected_max(const tools::wallet2 &w1, const cryptonote::subaddress_index exp_max)
 {
-  // these test assume we are starting with the default setup state
+  for (uint32_t i = 0; i <= exp_max.minor; ++i)
+  {
+    auto subaddr = w1.get_subaddress({exp_max.major, i});
+    EXPECT_NE(boost::none, w1.get_subaddress_index(subaddr));
+  }
+  auto subaddr = w1.get_subaddress({exp_max.major, exp_max.minor + 1});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(subaddr));
+};
+
+static void expect_default_wallet_state(const tools::wallet2 &w1)
+{
+  // these tests assume we are starting with the default setup state
   EXPECT_EQ(2, w1.get_num_subaddress_accounts());
   EXPECT_EQ(50, w1.get_subaddress_lookahead().first);
   EXPECT_EQ(200, w1.get_subaddress_lookahead().second);
+
+  // We assume we start with subaddrs for minor indexes 0 to 199
+  check_expected_max(w1, {0,199});
+  check_expected_max(w1, {1,199});
+  check_expected_max(w1, {49,199});
+  check_expected_max(w1, {50,199}); // 50 because the test starts with accounts 0 and 1 already allocated
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({51,0})));
+}
+
+TEST_F(WalletSubaddress, SetLookahead)
+{
+  expect_default_wallet_state(w1);
   // get_subaddress_index looks up keys in the private m_subaddresses dictionary so we will use it to test if a key is properly being scanned for
   cryptonote::subaddress_index test_idx = {50, 199};
   auto subaddr = w1.get_subaddress(test_idx);
@@ -117,14 +141,75 @@ TEST_F(WalletSubaddress, ExpandPubkeyTable)
   w1.set_subaddress_lookahead(100, 200);
   EXPECT_EQ(100, w1.get_subaddress_lookahead().first);
   EXPECT_EQ(200, w1.get_subaddress_lookahead().second);
-  test_idx = {100, 199};
-  subaddr = w1.get_subaddress(test_idx);
-  EXPECT_NE(boost::none, w1.get_subaddress_index(subaddr));
+  check_expected_max(w1, {100, 199});
   // next test expanding the minor lookahead
   w1.set_subaddress_lookahead(100, 300);
   EXPECT_EQ(100, w1.get_subaddress_lookahead().first);
   EXPECT_EQ(300, w1.get_subaddress_lookahead().second);
-  test_idx = {100, 299};
-  subaddr = w1.get_subaddress(test_idx);
-  EXPECT_NE(boost::none, w1.get_subaddress_index(subaddr));
+  check_expected_max(w1, {100, 299});
+}
+
+TEST_F(WalletSubaddress, ExpandThenSetMinorIncreaseOnly)
+{
+  expect_default_wallet_state(w1);
+
+  // Mock receive to {0,150}, so expand from there
+  w1.expand_subaddresses({0,150});
+  // We should now have subaddresses for minor indexes 0 to 349
+  check_expected_max(w1, {0,349});
+  check_expected_max(w1, {1,199});
+  check_expected_max(w1, {49,199});
+  check_expected_max(w1, {50,199});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({51,0})));
+
+  // Now set the minor lookahead 100 higher
+  w1.set_subaddress_lookahead(50, 200+100);
+  // We should have subaddresses for minor indexes 0 to 449
+  check_expected_max(w1, {0,449});
+  check_expected_max(w1, {1,299});
+  check_expected_max(w1, {49,299});
+  check_expected_max(w1, {50,299});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({51,0})));
+}
+
+TEST_F(WalletSubaddress, ExpandThenSetMajorIncreaseOnly)
+{
+  expect_default_wallet_state(w1);
+
+  // Mock receive to {40,0}, so expand from there
+  w1.expand_subaddresses({40,0});
+  check_expected_max(w1, {0,199});
+  check_expected_max(w1, {1,199});
+  check_expected_max(w1, {40,199});
+  check_expected_max(w1, {89,199});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({90,0})));
+
+  // Now set the major lookahead 10 higher
+  w1.set_subaddress_lookahead(50+10, 200);
+  check_expected_max(w1, {0,199});
+  check_expected_max(w1, {1,199});
+  check_expected_max(w1, {40,199});
+  check_expected_max(w1, {99,199});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({100,0})));
+}
+
+TEST_F(WalletSubaddress, ExpandThenSetIncreaseBoth)
+{
+  expect_default_wallet_state(w1);
+
+  // Mock receive to {40,150}, so expand from there
+  w1.expand_subaddresses({40,150});
+  check_expected_max(w1, {0,199});
+  check_expected_max(w1, {1,199});
+  check_expected_max(w1, {40,349});
+  check_expected_max(w1, {89,199});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({90,0})));
+
+  // Now set the major lookahead 10 higher and minor 100 higher
+  w1.set_subaddress_lookahead(50+10, 200+100);
+  check_expected_max(w1, {0,299});
+  check_expected_max(w1, {1,299});
+  check_expected_max(w1, {40,449});
+  check_expected_max(w1, {99,299});
+  EXPECT_EQ(boost::none, w1.get_subaddress_index(w1.get_subaddress({100,0})));
 }
diff --git a/utils/python-rpc/framework/wallet.py b/utils/python-rpc/framework/wallet.py
index d7cbde5da..d95c3b2b5 100644
--- a/utils/python-rpc/framework/wallet.py
+++ b/utils/python-rpc/framework/wallet.py
@@ -333,11 +333,12 @@ class Wallet(object):
         }
         return self.rpc.send_json_rpc_request(generate_from_keys)
 
-    def set_subaddress_lookahead(self, major_idx: int, minor_idx: int):
+    def set_subaddress_lookahead(self, major_idx: int, minor_idx: int, password = ""):
         lookahead = {
             'method': 'set_subaddress_lookahead',
             'jsonrpc': '2.0',
             'params' : {
+                'password': password,
                 'major_idx': major_idx,
                 'minor_idx': minor_idx
             },