From 6280b24dc41d0e12bf44e2c299d13291bef80961 Mon Sep 17 00:00:00 2001
From: j-berman <justinberman@protonmail.com>
Date: Thu, 28 Sep 2023 16:35:37 -0700
Subject: [PATCH] Monero: fix decoy selection algo and add test for latest
 spendable

- DSA only selected coinbase outputs and didn't match the wallet2
implementation
- Added test to make sure DSA will select a decoy output from the
most recent unlocked block
- Made usage of "height" in DSA consistent with other usage of
"height" in Monero code (height == num blocks in chain)
- Rely on monerod RPC response for output's unlocked status
---
 coins/monero/src/rpc/mod.rs                 | 68 ++++++++---------
 coins/monero/src/transaction.rs             |  2 +
 coins/monero/src/wallet/decoys.rs           | 43 ++++++++---
 coins/monero/tests/decoys.rs                | 85 +++++++++++++++++++++
 coins/monero/tests/runner.rs                |  2 +-
 coins/monero/tests/send.rs                  |  2 +-
 processor/src/networks/monero.rs            |  2 +-
 tests/full-stack/src/tests/mint_and_burn.rs |  2 +-
 tests/processor/src/networks.rs             |  2 +-
 9 files changed, 155 insertions(+), 53 deletions(-)
 create mode 100644 coins/monero/tests/decoys.rs

diff --git a/coins/monero/src/rpc/mod.rs b/coins/monero/src/rpc/mod.rs
index cfe4e4b8d..35b13daa0 100644
--- a/coins/monero/src/rpc/mod.rs
+++ b/coins/monero/src/rpc/mod.rs
@@ -17,7 +17,7 @@ use serde_json::{Value, json};
 use crate::{
   Protocol,
   serialize::*,
-  transaction::{Input, Timelock, Transaction},
+  transaction::{Input, Transaction},
   block::Block,
   wallet::{FeePriority, Fee},
 };
@@ -52,6 +52,14 @@ struct TransactionsResponse {
   txs: Vec<TransactionResponse>,
 }
 
+#[derive(Deserialize, Debug)]
+pub struct OutputResponse {
+  pub height: usize,
+  pub unlocked: bool,
+  key: String,
+  mask: String,
+}
+
 #[derive(Clone, PartialEq, Eq, Debug)]
 #[cfg_attr(feature = "std", derive(thiserror::Error))]
 pub enum RpcError {
@@ -532,29 +540,15 @@ impl<R: RpcConnection> Rpc<R> {
     Ok(distributions.distributions.swap_remove(0).distribution)
   }
 
-  /// Get the specified outputs from the RingCT (zero-amount) pool, but only return them if their
-  /// timelock has been satisfied.
-  ///
-  /// The timelock being satisfied is distinct from being free of the 10-block lock applied to all
-  /// Monero transactions.
-  pub async fn get_unlocked_outputs(
-    &self,
-    indexes: &[u64],
-    height: usize,
-  ) -> Result<Vec<Option<[EdwardsPoint; 2]>>, RpcError> {
-    #[derive(Deserialize, Debug)]
-    struct Out {
-      key: String,
-      mask: String,
-      txid: String,
-    }
-
+  /// Get the specified outputs from the RingCT (zero-amount) pool
+  pub async fn get_outs(&self, indexes: &[u64]) -> Result<Vec<OutputResponse>, RpcError> {
     #[derive(Deserialize, Debug)]
-    struct Outs {
-      outs: Vec<Out>,
+    struct OutsResponse {
+      status: String,
+      outs: Vec<OutputResponse>,
     }
 
-    let outs: Outs = self
+    let res: OutsResponse = self
       .rpc_call(
         "get_outs",
         Some(json!({
@@ -567,22 +561,25 @@ impl<R: RpcConnection> Rpc<R> {
       )
       .await?;
 
-    let txs = self
-      .get_transactions(
-        &outs
-          .outs
-          .iter()
-          .map(|out| rpc_hex(&out.txid)?.try_into().map_err(|_| RpcError::InvalidNode))
-          .collect::<Result<Vec<_>, _>>()?,
-      )
-      .await?;
+    if res.status != "OK" {
+      Err(RpcError::InvalidNode)?;
+    }
+
+    Ok(res.outs)
+  }
+
+  /// Get the specified outputs from the RingCT (zero-amount) pool, but only return them
+  /// if they are unlocked.
+  pub async fn get_unlocked_outputs(
+    &self,
+    indexes: &[u64],
+  ) -> Result<Vec<Option<[EdwardsPoint; 2]>>, RpcError> {
+    let outs: Vec<OutputResponse> = self.get_outs(indexes).await?;
 
-    // TODO: https://github.com/serai-dex/serai/issues/104
     outs
-      .outs
       .iter()
       .enumerate()
-      .map(|(i, out)| {
+      .map(|(_, out)| {
         // Allow keys to be invalid, though if they are, return None to trigger selection of a new
         // decoy
         // Only valid keys can be used in CLSAG proofs, hence the need for re-selection, yet
@@ -597,10 +594,7 @@ impl<R: RpcConnection> Rpc<R> {
         .decompress() else {
           return Ok(None);
         };
-        Ok(
-          Some([key, rpc_point(&out.mask)?])
-            .filter(|_| Timelock::Block(height) >= txs[i].prefix.timelock),
-        )
+        Ok(Some([key, rpc_point(&out.mask)?]).filter(|_| out.unlocked))
       })
       .collect()
   }
diff --git a/coins/monero/src/transaction.rs b/coins/monero/src/transaction.rs
index 5902bc2f1..51764d288 100644
--- a/coins/monero/src/transaction.rs
+++ b/coins/monero/src/transaction.rs
@@ -168,7 +168,9 @@ impl Timelock {
 impl PartialOrd for Timelock {
   fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
     match (self, other) {
+      (Timelock::None, Timelock::None) => Some(Ordering::Equal),
       (Timelock::None, _) => Some(Ordering::Less),
+      (_, Timelock::None) => Some(Ordering::Greater),
       (Timelock::Block(a), Timelock::Block(b)) => a.partial_cmp(b),
       (Timelock::Time(a), Timelock::Time(b)) => a.partial_cmp(b),
       _ => None,
diff --git a/coins/monero/src/wallet/decoys.rs b/coins/monero/src/wallet/decoys.rs
index fd1b8bcde..064bae06a 100644
--- a/coins/monero/src/wallet/decoys.rs
+++ b/coins/monero/src/wallet/decoys.rs
@@ -20,8 +20,8 @@ use crate::{
   rpc::{RpcError, RpcConnection, Rpc},
 };
 
-const LOCK_WINDOW: usize = 10;
-const MATURITY: u64 = 60;
+pub const LOCK_WINDOW: usize = 10;
+const COINBASE_LOCK_WINDOW: u64 = 60;
 const RECENT_WINDOW: usize = 15;
 const BLOCK_TIME: usize = 120;
 const BLOCKS_PER_YEAR: usize = 365 * 24 * 60 * 60 / BLOCK_TIME;
@@ -47,7 +47,7 @@ async fn select_n<'a, R: RngCore + CryptoRng, RPC: RpcConnection>(
   used: &mut HashSet<u64>,
   count: usize,
 ) -> Result<Vec<(u64, [EdwardsPoint; 2])>, RpcError> {
-  if height >= rpc.get_height().await? {
+  if height > rpc.get_height().await? {
     // TODO: Don't use InternalError for the caller's failure
     Err(RpcError::InternalError("decoys being requested from too young blocks"))?;
   }
@@ -58,6 +58,8 @@ async fn select_n<'a, R: RngCore + CryptoRng, RPC: RpcConnection>(
   // Retries on failure. Retries are obvious as decoys, yet should be minimal
   while confirmed.len() != count {
     let remaining = count - confirmed.len();
+    // TODO: over-request candidates in case some are locked to avoid needing
+    // round trips to the daemon (and revealing obvious decoys to the daemon)
     let mut candidates = Vec::with_capacity(remaining);
     while candidates.len() != remaining {
       #[cfg(test)]
@@ -109,7 +111,9 @@ async fn select_n<'a, R: RngCore + CryptoRng, RPC: RpcConnection>(
       }
     }
 
-    for (i, output) in rpc.get_unlocked_outputs(&candidates, height).await?.iter_mut().enumerate() {
+    // TODO: make sure that the real output is included in the response, and
+    // that mask and key are equal to expected
+    for (i, output) in rpc.get_unlocked_outputs(&candidates).await?.iter_mut().enumerate() {
       // Don't include the real spend as a decoy, despite requesting it
       if real_indexes.contains(&i) {
         continue;
@@ -151,6 +155,14 @@ impl Decoys {
     self.offsets.len()
   }
 
+  pub fn indexes(&self) -> Vec<u64> {
+    let mut res = vec![self.offsets[0]; self.len()];
+    for m in 1 .. res.len() {
+      res[m] = res[m - 1] + self.offsets[m];
+    }
+    res
+  }
+
   /// Select decoys using the same distribution as Monero.
   pub async fn select<R: RngCore + CryptoRng, RPC: RpcConnection>(
     rng: &mut R,
@@ -174,18 +186,26 @@ impl Decoys {
       outputs.push((real[real.len() - 1], [input.key(), input.commitment().calculate()]));
     }
 
-    if distribution.len() <= height {
-      let extension = rpc.get_output_distribution(distribution.len(), height).await?;
+    if distribution.len() < height {
+      // TODO: verify distribution elems are strictly increasing
+      let extension =
+        rpc.get_output_distribution(distribution.len(), height.saturating_sub(1)).await?;
       distribution.extend(extension);
     }
     // If asked to use an older height than previously asked, truncate to ensure accuracy
     // Should never happen, yet risks desyncing if it did
-    distribution.truncate(height + 1); // height is inclusive, and 0 is a valid height
+    distribution.truncate(height);
+
+    if distribution.len() != height {
+      Err(RpcError::InternalError("unexpected rct out distribution len"))?;
+    } else if distribution.len() < LOCK_WINDOW {
+      Err(RpcError::InternalError("not enough decoy candidates"))?;
+    }
 
-    let high = distribution[distribution.len() - 1];
     let per_second = {
       let blocks = distribution.len().min(BLOCKS_PER_YEAR);
-      let outputs = high - distribution[distribution.len().saturating_sub(blocks + 1)];
+      let initial = distribution[distribution.len().saturating_sub(blocks + 1)];
+      let outputs = distribution[distribution.len() - 1].saturating_sub(initial);
       (outputs as f64) / ((blocks * BLOCK_TIME) as f64)
     };
 
@@ -195,8 +215,9 @@ impl Decoys {
     }
 
     // TODO: Create a TX with less than the target amount, as allowed by the protocol
-    if (high - MATURITY) < u64::try_from(inputs.len() * ring_len).unwrap() {
-      Err(RpcError::InternalError("not enough decoy candidates"))?;
+    let high = distribution[distribution.len() - LOCK_WINDOW];
+    if high.saturating_sub(COINBASE_LOCK_WINDOW) < u64::try_from(inputs.len() * ring_len).unwrap() {
+      Err(RpcError::InternalError("not enough coinbase candidates"))?;
     }
 
     // Select all decoys for this transaction, assuming we generate a sane transaction
diff --git a/coins/monero/tests/decoys.rs b/coins/monero/tests/decoys.rs
new file mode 100644
index 000000000..0bb05e707
--- /dev/null
+++ b/coins/monero/tests/decoys.rs
@@ -0,0 +1,85 @@
+use monero_serai::{
+  transaction::Transaction,
+  wallet::{SpendableOutput, decoys::LOCK_WINDOW},
+  rpc::{Rpc, OutputResponse},
+  Protocol,
+};
+
+mod runner;
+
+test!(
+  select_latest_output_as_decoy,
+  (
+    // First make an initial tx0
+    |_, mut builder: Builder, addr| async move {
+      builder.add_payment(addr, 2000000000000);
+      (builder.build().unwrap(), ())
+    },
+    |rpc: Rpc<_>, tx: Transaction, mut scanner: Scanner, _| async move {
+      let output = scanner.scan_transaction(&tx).not_locked().swap_remove(0);
+      assert_eq!(output.commitment().amount, 2000000000000);
+      SpendableOutput::from(&rpc, output).await.unwrap()
+    },
+  ),
+  (
+    // Then make a second tx1
+    |protocol: Protocol, rpc: Rpc<_>, mut builder: Builder, addr, state: _| async move {
+      let output_tx0: SpendableOutput = state;
+      let decoys = Decoys::select(
+        &mut OsRng,
+        &rpc,
+        protocol.ring_len(),
+        rpc.get_height().await.unwrap(),
+        &[output_tx0.clone()],
+      )
+      .await
+      .unwrap();
+
+      let inputs = [output_tx0.clone()].into_iter().zip(decoys).collect::<Vec<_>>();
+      builder.add_inputs(&inputs);
+      builder.add_payment(addr, 1000000000000);
+
+      (builder.build().unwrap(), (protocol, output_tx0))
+    },
+    // Then make sure DSA selects freshly unlocked output from tx1 as a decoy
+    |rpc: Rpc<_>, tx: Transaction, mut scanner: Scanner, state: (_, _)| async move {
+      use rand_core::OsRng;
+
+      let height = rpc.get_height().await.unwrap();
+
+      let output_tx1 =
+        SpendableOutput::from(&rpc, scanner.scan_transaction(&tx).not_locked().swap_remove(0))
+          .await
+          .unwrap();
+
+      // Make sure output from tx1 is in the block in which it unlocks
+      let out_tx1: OutputResponse =
+        rpc.get_outs(&[output_tx1.global_index]).await.unwrap().swap_remove(0);
+      assert_eq!(out_tx1.height, height - LOCK_WINDOW);
+      assert!(out_tx1.unlocked);
+
+      // Select decoys using spendable output from tx0 as the real, and make sure DSA selects
+      // the freshly unlocked output from tx1 as a decoy
+      let (protocol, output_tx0): (Protocol, SpendableOutput) = state;
+      let mut selected_fresh_decoy = false;
+      let mut attempts = 1000;
+      while !selected_fresh_decoy && attempts > 0 {
+        let decoys = Decoys::select(
+          &mut OsRng, // TODO: use a seeded RNG to consistently select the latest output
+          &rpc,
+          protocol.ring_len(),
+          height,
+          &[output_tx0.clone()],
+        )
+        .await
+        .unwrap();
+
+        selected_fresh_decoy = decoys[0].indexes().contains(&output_tx1.global_index);
+        attempts -= 1;
+      }
+
+      assert!(selected_fresh_decoy);
+      assert_eq!(height, rpc.get_height().await.unwrap());
+    },
+  ),
+);
diff --git a/coins/monero/tests/runner.rs b/coins/monero/tests/runner.rs
index 9362d8605..5aa8b08c8 100644
--- a/coins/monero/tests/runner.rs
+++ b/coins/monero/tests/runner.rs
@@ -264,7 +264,7 @@ macro_rules! test {
               &mut OsRng,
               &rpc,
               protocol.ring_len(),
-              rpc.get_height().await.unwrap() - 1,
+              rpc.get_height().await.unwrap(),
               &[miner_tx.clone()]
             )
             .await
diff --git a/coins/monero/tests/send.rs b/coins/monero/tests/send.rs
index b2abcb67b..d08b780a0 100644
--- a/coins/monero/tests/send.rs
+++ b/coins/monero/tests/send.rs
@@ -28,7 +28,7 @@ async fn add_inputs(
     &mut OsRng,
     rpc,
     protocol.ring_len(),
-    rpc.get_height().await.unwrap() - 1,
+    rpc.get_height().await.unwrap(),
     &spendable_outputs,
   )
   .await
diff --git a/processor/src/networks/monero.rs b/processor/src/networks/monero.rs
index 82f04326e..4aff1e0fb 100644
--- a/processor/src/networks/monero.rs
+++ b/processor/src/networks/monero.rs
@@ -631,7 +631,7 @@ impl Network for Monero {
       &mut OsRng,
       &self.rpc,
       protocol.ring_len(),
-      self.rpc.get_height().await.unwrap() - 1,
+      self.rpc.get_height().await.unwrap(),
       &outputs,
     )
     .await
diff --git a/tests/full-stack/src/tests/mint_and_burn.rs b/tests/full-stack/src/tests/mint_and_burn.rs
index b69f03244..2aa8f761e 100644
--- a/tests/full-stack/src/tests/mint_and_burn.rs
+++ b/tests/full-stack/src/tests/mint_and_burn.rs
@@ -356,7 +356,7 @@ async fn mint_and_burn_test() {
           &mut OsRng,
           &rpc,
           Protocol::v16.ring_len(),
-          rpc.get_height().await.unwrap() - 1,
+          rpc.get_height().await.unwrap(),
           &[output.clone()],
         )
         .await
diff --git a/tests/processor/src/networks.rs b/tests/processor/src/networks.rs
index d2b7cd4b2..9a23286d3 100644
--- a/tests/processor/src/networks.rs
+++ b/tests/processor/src/networks.rs
@@ -332,7 +332,7 @@ impl Wallet {
           &mut OsRng,
           &rpc,
           Protocol::v16.ring_len(),
-          rpc.get_height().await.unwrap() - 1,
+          rpc.get_height().await.unwrap(),
           &these_inputs,
         )
         .await