From 9ebf438645505472f8134522dbc2d6805bb2a488 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Tue, 4 Jul 2023 11:59:50 -0400
Subject: [PATCH] Add necessary checks to Eventuality re: supported protocols

---
 coins/monero/src/lib.rs             | 26 +++++++++++++++++++++-----
 coins/monero/src/ringct/mod.rs      | 10 +++++-----
 coins/monero/src/wallet/send/mod.rs | 10 ++++++++++
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/coins/monero/src/lib.rs b/coins/monero/src/lib.rs
index cd5b4d96..74b42f81 100644
--- a/coins/monero/src/lib.rs
+++ b/coins/monero/src/lib.rs
@@ -25,6 +25,7 @@ use serialize::{read_byte, read_u16};
 
 /// RingCT structs and functionality.
 pub mod ringct;
+use ringct::RctType;
 
 /// Transaction structs.
 pub mod transaction;
@@ -45,14 +46,16 @@ pub(crate) fn INV_EIGHT() -> Scalar {
   *INV_EIGHT_CELL.get_or_init(|| Scalar::from(8u8).invert())
 }
 
-/// Monero protocol version. v15 is omitted as v15 was simply v14 and v16 being active at the same
-/// time, with regards to the transactions supported. Accordingly, v16 should be used during v15.
+/// Monero protocol version.
+///
+/// v15 is omitted as v15 was simply v14 and v16 being active at the same time, with regards to the
+/// transactions supported. Accordingly, v16 should be used during v15.
 #[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize)]
 #[allow(non_camel_case_types)]
 pub enum Protocol {
   v14,
   v16,
-  Custom { ring_len: usize, bp_plus: bool },
+  Custom { ring_len: usize, bp_plus: bool, optimal_rct_type: RctType },
 }
 
 impl Protocol {
@@ -66,6 +69,7 @@ impl Protocol {
   }
 
   /// Whether or not the specified version uses Bulletproofs or Bulletproofs+.
+  ///
   /// This method will likely be reworked when versions not using Bulletproofs at all are added.
   pub fn bp_plus(&self) -> bool {
     match self {
@@ -75,15 +79,25 @@ impl Protocol {
     }
   }
 
+  // TODO: Make this an Option when we support pre-RCT protocols
+  pub fn optimal_rct_type(&self) -> RctType {
+    match self {
+      Protocol::v14 => RctType::Clsag,
+      Protocol::v16 => RctType::BulletproofsPlus,
+      Protocol::Custom { optimal_rct_type, .. } => *optimal_rct_type,
+    }
+  }
+
   pub(crate) fn write<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
     match self {
       Protocol::v14 => w.write_all(&[0, 14]),
       Protocol::v16 => w.write_all(&[0, 16]),
-      Protocol::Custom { ring_len, bp_plus } => {
+      Protocol::Custom { ring_len, bp_plus, optimal_rct_type } => {
         // Custom, version 0
         w.write_all(&[1, 0])?;
         w.write_all(&u16::try_from(*ring_len).unwrap().to_le_bytes())?;
-        w.write_all(&[u8::from(*bp_plus)])
+        w.write_all(&[u8::from(*bp_plus)])?;
+        w.write_all(&[optimal_rct_type.to_byte()])
       }
     }
   }
@@ -105,6 +119,8 @@ impl Protocol {
             1 => true,
             _ => Err(io::Error::new(io::ErrorKind::Other, "invalid bool serialization"))?,
           },
+          optimal_rct_type: RctType::from_byte(read_byte(r)?)
+            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "invalid RctType serialization"))?,
         },
         _ => {
           Err(io::Error::new(io::ErrorKind::Other, "unrecognized custom protocol serialization"))?
diff --git a/coins/monero/src/ringct/mod.rs b/coins/monero/src/ringct/mod.rs
index 1862e9e5..a706e3e4 100644
--- a/coins/monero/src/ringct/mod.rs
+++ b/coins/monero/src/ringct/mod.rs
@@ -4,7 +4,7 @@ use std_shims::{
   io::{self, Read, Write},
 };
 
-use zeroize::Zeroizing;
+use zeroize::{Zeroize, Zeroizing};
 
 use curve25519_dalek::{constants::ED25519_BASEPOINT_TABLE, scalar::Scalar, edwards::EdwardsPoint};
 
@@ -57,7 +57,7 @@ impl EncryptedAmount {
   }
 }
 
-#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+#[derive(Clone, Copy, PartialEq, Eq, Debug, Zeroize)]
 pub enum RctType {
   /// No RCT proofs.
   Null,
@@ -77,7 +77,7 @@ pub enum RctType {
 }
 
 impl RctType {
-  fn to_byte(self) -> u8 {
+  pub fn to_byte(self) -> u8 {
     match self {
       RctType::Null => 0,
       RctType::MlsagAggregate => 1,
@@ -89,7 +89,7 @@ impl RctType {
     }
   }
 
-  fn from_byte(byte: u8) -> Option<Self> {
+  pub fn from_byte(byte: u8) -> Option<Self> {
     Some(match byte {
       0 => RctType::Null,
       1 => RctType::MlsagAggregate,
@@ -102,7 +102,7 @@ impl RctType {
     })
   }
 
-  fn compact_encrypted_amounts(&self) -> bool {
+  pub fn compact_encrypted_amounts(&self) -> bool {
     match self {
       RctType::Null => false,
       RctType::MlsagAggregate => false,
diff --git a/coins/monero/src/wallet/send/mod.rs b/coins/monero/src/wallet/send/mod.rs
index 497e8974..7eb1fcf2 100644
--- a/coins/monero/src/wallet/send/mod.rs
+++ b/coins/monero/src/wallet/send/mod.rs
@@ -746,6 +746,16 @@ impl Eventuality {
       uniqueness(&tx.prefix.inputs),
     );
 
+    let rct_type = tx.rct_signatures.rct_type();
+    if rct_type != self.protocol.optimal_rct_type() {
+      return false;
+    }
+
+    // TODO: Remove this when the following for loop is updated
+    if !rct_type.compact_encrypted_amounts() {
+      panic!("created an Eventuality for a very old RctType we don't support proving for");
+    }
+
     for (o, (expected, actual)) in outputs.iter().zip(tx.prefix.outputs.iter()).enumerate() {
       // Verify the output, commitment, and encrypted amount.
       if (&Output {