wallet: add a set_ring command

This is so one can set rings for spent key images in case the
attackers don't merge the ring matching patch set.

6 files changed, 124 insertions, 1 deletions

diff --git a/src/simplewallet/simplewallet.cpp b/src/simplewallet/simplewallet.cpp
index 89edee3fe..61a234317 100644
--- a/src/simplewallet/simplewallet.cpp
+++ b/src/simplewallet/simplewallet.cpp
@@ -1317,7 +1317,7 @@ bool simple_wallet::print_ring(const std::vector<std::string> &args)
       std::stringstream str;
       for (const auto &x: ring)
         str << x << " ";
-      success_msg_writer() << tr("Ring size ") << std::to_string(ring.size()) << ": " << str.str();
+      success_msg_writer() << tr("Ring size ") << std::to_string(ring.size()) << ": " << str.str() << tr(" (absolute)");
     }
     else
     {
@@ -1332,6 +1332,81 @@ bool simple_wallet::print_ring(const std::vector<std::string> &args)
   return true;
 }
 
+bool simple_wallet::set_ring(const std::vector<std::string> &args)
+{
+  crypto::key_image key_image;
+  if (args.size() < 3)
+  {
+    fail_msg_writer() << tr("usage: set_ring <key_image> absolute|relative <index> [<index>...]");
+    return true;
+  }
+
+  if (!epee::string_tools::hex_to_pod(args[0], key_image))
+  {
+    fail_msg_writer() << tr("Invalid key image");
+    return true;
+  }
+
+  bool relative;
+  if (args[1] == "absolute")
+  {
+    relative = false;
+  }
+  else if (args[1] == "relative")
+  {
+    relative = true;
+  }
+  else
+  {
+    fail_msg_writer() << tr("Missing absolute or relative keyword");
+    return true;
+  }
+
+  std::vector<uint64_t> ring;
+  for (size_t n = 2; n < args.size(); ++n)
+  {
+    ring.resize(ring.size() + 1);
+    if (!string_tools::get_xtype_from_string(ring.back(), args[n]))
+    {
+      fail_msg_writer() << tr("invalid index: must be a strictly positive unsigned integer");
+      return true;
+    }
+    if (relative)
+    {
+      if (ring.size() > 1 && !ring.back())
+      {
+        fail_msg_writer() << tr("invalid index: must be a strictly positive unsigned integer");
+        return true;
+      }
+      uint64_t sum = 0;
+      for (uint64_t out: ring)
+      {
+        if (out > std::numeric_limits<uint64_t>::max() - sum)
+        {
+          fail_msg_writer() << tr("invalid index: indices wrap");
+          return true;
+        }
+        sum += out;
+      }
+    }
+    else
+    {
+      if (ring.size() > 1 && ring[ring.size() - 2] >= ring[ring.size() - 1])
+      {
+        fail_msg_writer() << tr("invalid index: indices should be in strictly ascending order");
+        return true;
+      }
+    }
+  }
+  if (!m_wallet->set_ring(key_image, ring, relative))
+  {
+    fail_msg_writer() << tr("failed to set ring");
+    return true;
+  }
+
+  return true;
+}
+
 bool simple_wallet::blackball(const std::vector<std::string> &args)
 {
   crypto::public_key output;
@@ -2166,6 +2241,10 @@ simple_wallet::simple_wallet()
                            boost::bind(&simple_wallet::print_ring, this, _1),
                            tr("print_ring <key_image>"),
                            tr("Print the ring used to spend a given key image (if the ring size is > 1)"));
+  m_cmd_binder.set_handler("set_ring",
+                           boost::bind(&simple_wallet::set_ring, this, _1),
+                           tr("set_ring <key_image> absolute|relative <index> [<index>...]"),
+                           tr("Set the ring used for a given key image, so it can be reused in a fork"));
   m_cmd_binder.set_handler("save_known_rings",
                            boost::bind(&simple_wallet::save_known_rings, this, _1),
                            tr("save_known_rings"),
diff --git a/src/simplewallet/simplewallet.h b/src/simplewallet/simplewallet.h
index 372057022..c69df817d 100644
--- a/src/simplewallet/simplewallet.h
+++ b/src/simplewallet/simplewallet.h
@@ -211,6 +211,7 @@ namespace cryptonote
     bool submit_multisig(const std::vector<std::string>& args);
     bool export_raw_multisig(const std::vector<std::string>& args);
     bool print_ring(const std::vector<std::string>& args);
+    bool set_ring(const std::vector<std::string>& args);
     bool save_known_rings(const std::vector<std::string>& args);
     bool blackball(const std::vector<std::string>& args);
     bool unblackball(const std::vector<std::string>& args);
diff --git a/src/wallet/ringdb.cpp b/src/wallet/ringdb.cpp
index bd2b4453b..d2ed90e51 100644
--- a/src/wallet/ringdb.cpp
+++ b/src/wallet/ringdb.cpp
@@ -340,6 +340,36 @@ bool ringdb::get_ring(const crypto::chacha_key &chacha_key, const crypto::key_im
   return true;
 }
 
+bool ringdb::set_ring(const crypto::chacha_key &chacha_key, const crypto::key_image &key_image, const std::vector<uint64_t> &outs, bool relative)
+{
+  MDB_txn *txn;
+  int dbr;
+  bool tx_active = false;
+
+  dbr = resize_env(env, filename.c_str(), outs.size() * 64);
+  THROW_WALLET_EXCEPTION_IF(dbr, tools::error::wallet_internal_error, "Failed to set env map size: " + std::string(mdb_strerror(dbr)));
+  dbr = mdb_txn_begin(env, NULL, 0, &txn);
+  THROW_WALLET_EXCEPTION_IF(dbr, tools::error::wallet_internal_error, "Failed to create LMDB transaction: " + std::string(mdb_strerror(dbr)));
+  epee::misc_utils::auto_scope_leave_caller txn_dtor = epee::misc_utils::create_scope_leave_handler([&](){if (tx_active) mdb_txn_abort(txn);});
+  tx_active = true;
+
+  MDB_val key, data;
+  std::string key_ciphertext = encrypt(key_image, chacha_key);
+  key.mv_data = (void*)key_ciphertext.data();
+  key.mv_size = key_ciphertext.size();
+  std::string compressed_ring = compress_ring(relative ? outs : cryptonote::absolute_output_offsets_to_relative(outs));
+  std::string data_ciphertext = encrypt(compressed_ring, key_image, chacha_key);
+  data.mv_size = data_ciphertext.size();
+  data.mv_data = (void*)data_ciphertext.c_str();
+  dbr = mdb_put(txn, dbi_rings, &key, &data, 0);
+  THROW_WALLET_EXCEPTION_IF(dbr, tools::error::wallet_internal_error, "Failed to set ring for key image in LMDB table: " + std::string(mdb_strerror(dbr)));
+
+  dbr = mdb_txn_commit(txn);
+  THROW_WALLET_EXCEPTION_IF(dbr, tools::error::wallet_internal_error, "Failed to commit txn setting ring to database: " + std::string(mdb_strerror(dbr)));
+  tx_active = false;
+  return true;
+}
+
 bool ringdb::blackball_worker(const crypto::public_key &output, int op)
 {
   MDB_txn *txn;
diff --git a/src/wallet/ringdb.h b/src/wallet/ringdb.h
index 351ae5a2b..fb6b732a9 100644
--- a/src/wallet/ringdb.h
+++ b/src/wallet/ringdb.h
@@ -46,6 +46,7 @@ namespace tools
     bool add_rings(const crypto::chacha_key &chacha_key, const cryptonote::transaction_prefix &tx);
     bool remove_rings(const crypto::chacha_key &chacha_key, const cryptonote::transaction_prefix &tx);
     bool get_ring(const crypto::chacha_key &chacha_key, const crypto::key_image &key_image, std::vector<uint64_t> &outs);
+    bool set_ring(const crypto::chacha_key &chacha_key, const crypto::key_image &key_image, const std::vector<uint64_t> &outs, bool relative);
 
     bool blackball(const crypto::public_key &output);
     bool unblackball(const crypto::public_key &output);
diff --git a/src/wallet/wallet2.cpp b/src/wallet/wallet2.cpp
index 732dc0789..b2cbd416b 100644
--- a/src/wallet/wallet2.cpp
+++ b/src/wallet/wallet2.cpp
@@ -5496,6 +5496,17 @@ bool wallet2::get_ring(const crypto::key_image &key_image, std::vector<uint64_t>
   return get_ring(key, key_image, outs);
 }
 
+bool wallet2::set_ring(const crypto::key_image &key_image, const std::vector<uint64_t> &outs, bool relative)
+{
+  if (!m_ringdb)
+    return true;
+
+  crypto::chacha_key key;
+  generate_chacha_key_from_secret_keys(key);
+
+  return m_ringdb->set_ring(key, key_image, outs, relative);
+}
+
 bool wallet2::find_and_save_rings(bool force)
 {
   if (!force && m_ring_history_saved)
diff --git a/src/wallet/wallet2.h b/src/wallet/wallet2.h
index c62bdbfc8..acccc6ca8 100644
--- a/src/wallet/wallet2.h
+++ b/src/wallet/wallet2.h
@@ -1059,6 +1059,7 @@ namespace tools
     void set_ring_database(const std::string &filename);
     const std::string get_ring_database() const { return m_ring_database; }
     bool get_ring(const crypto::key_image &key_image, std::vector<uint64_t> &outs);
+    bool set_ring(const crypto::key_image &key_image, const std::vector<uint64_t> &outs, bool relative);
     bool find_and_save_rings(bool force = true);
 
     bool blackball_output(const crypto::public_key &output);