From 534e1bb11d22b126aa1816a26530ef9af96206a7 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Sun, 26 Mar 2023 03:43:51 -0400
Subject: [PATCH] Fix Monero's Extra::fee_weight and handling of data limits

---
 coins/monero/src/wallet/extra.rs    | 13 +++++++++----
 coins/monero/src/wallet/send/mod.rs |  5 +++--
 coins/monero/tests/add_data.rs      | 13 ++++++++-----
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/coins/monero/src/wallet/extra.rs b/coins/monero/src/wallet/extra.rs
index 1b902109..d5330f84 100644
--- a/coins/monero/src/wallet/extra.rs
+++ b/coins/monero/src/wallet/extra.rs
@@ -172,15 +172,20 @@ impl Extra {
   }
 
   #[rustfmt::skip]
-  pub(crate) fn fee_weight(outputs: usize, payment_id: bool, data: &[Vec<u8>]) -> usize {
+  pub(crate) fn fee_weight(
+    outputs: usize,
+    additional: bool,
+    payment_id: bool,
+    data: &[Vec<u8>]
+  ) -> usize {
     // PublicKey, key
     (1 + 32) +
     // PublicKeys, length, additional keys
-    (1 + 1 + (outputs.saturating_sub(1) * 32)) +
+    (if additional { 1 + 1 + (outputs * 32) } else { 0 }) +
     // PaymentId (Nonce), length, encrypted, ID
     (if payment_id { 1 + 1 + 1 + 8 } else { 0 }) +
-    // Nonce, length, data (if existent)
-    data.iter().map(|v| 1 + varint_len(v.len()) + v.len()).sum::<usize>()
+    // Nonce, length, ARBITRARY_DATA_MARKER, data
+    data.iter().map(|v| 1 + varint_len(1 + v.len()) + 1 + v.len()).sum::<usize>()
   }
 
   pub fn write<W: Write>(&self, w: &mut W) -> io::Result<()> {
diff --git a/coins/monero/src/wallet/send/mod.rs b/coins/monero/src/wallet/send/mod.rs
index 5b9c9320..23c5702b 100644
--- a/coins/monero/src/wallet/send/mod.rs
+++ b/coins/monero/src/wallet/send/mod.rs
@@ -337,7 +337,8 @@ impl SignableTransaction {
     has_payment_id |= outputs == 2;
 
     // Calculate the extra length
-    let extra = Extra::fee_weight(outputs, has_payment_id, data.as_ref());
+    // Assume additional keys are needed in order to cause a worst-case estimation
+    let extra = Extra::fee_weight(outputs, true, has_payment_id, data.as_ref());
 
     // https://github.com/monero-project/monero/pull/8733
     const MAX_EXTRA_SIZE: usize = 1060;
@@ -541,7 +542,7 @@ impl SignableTransaction {
     }
 
     // Include data if present
-    let extra_len = Extra::fee_weight(Rs_len, id.is_some(), data.as_ref());
+    let extra_len = Extra::fee_weight(Rs_len, additional, id.is_some(), data.as_ref());
     for part in data.drain(..) {
       let mut arb = vec![ARBITRARY_DATA_MARKER];
       arb.extend(part);
diff --git a/coins/monero/tests/add_data.rs b/coins/monero/tests/add_data.rs
index fe967c72..9f53b72d 100644
--- a/coins/monero/tests/add_data.rs
+++ b/coins/monero/tests/add_data.rs
@@ -30,10 +30,13 @@ test!(
   add_multiple_data_less_than_max,
   (
     |_, mut builder: Builder, addr| async move {
-      let data = vec![b'\0'; MAX_ARBITRARY_DATA_SIZE - 1];
+      let mut data = vec![];
+      for b in 1 ..= 3 {
+        data.push(vec![b; MAX_ARBITRARY_DATA_SIZE - 1]);
+      }
 
-      // Add tx multiple times
-      for _ in 0 .. 5 {
+      // Add data multiple times
+      for data in &data {
         let result = builder.add_data(data.clone());
         assert!(result.is_ok());
       }
@@ -41,10 +44,10 @@ test!(
       builder.add_payment(addr, 5);
       (builder.build().unwrap(), data)
     },
-    |_, tx: Transaction, mut scanner: Scanner, data: Vec<u8>| async move {
+    |_, tx: Transaction, mut scanner: Scanner, data: Vec<Vec<u8>>| async move {
       let output = scanner.scan_transaction(&tx).not_locked().swap_remove(0);
       assert_eq!(output.commitment().amount, 5);
-      assert_eq!(output.arbitrary_data(), vec![data; 5]);
+      assert_eq!(output.arbitrary_data(), data);
     },
   ),
 );