diff --git a/transaction/builder/src/memo_builder/rth_memo_builder.rs b/transaction/builder/src/memo_builder/rth_memo_builder.rs index 3b1bac0010..0d5e03e1a9 100644 --- a/transaction/builder/src/memo_builder/rth_memo_builder.rs +++ b/transaction/builder/src/memo_builder/rth_memo_builder.rs @@ -6,10 +6,10 @@ //! (RTH is an abbrevation of Recoverable Transaction History.) //! This MemoBuilder policy implements Recoverable Transaction History using //! the encrypted memos, as envisioned in MCIP #4. - use super::MemoBuilder; use crate::ReservedSubaddresses; -use alloc::{boxed::Box, fmt::Debug, format, sync::Arc}; +use alloc::{boxed::Box, fmt::Debug, format, string::String, sync::Arc}; +use displaydoc::Display; use mc_account_keys::{PublicAddress, ShortAddressHash}; use mc_transaction_core::{ tokens::Mob, Amount, MemoContext, MemoPayload, NewMemoError, Token, TokenId, @@ -59,12 +59,8 @@ use mc_transaction_extra::{ pub struct RTHMemoBuilder { // The credential used to form 0x0100 and 0x0101 memos, if present. sender_cred: Option, - // The payment request id, if any - payment_request_id: Option, - // The payment intent id, if any - payment_intent_id: Option, - // The memo generator function, if any - flexible_memo_generator: Option, + // Different options for the custom memo data + custom_memo_type: Option, // Whether destination memos are enabled. destination_memo_enabled: bool, // Tracks if we already wrote a destination memo, for error reporting @@ -81,6 +77,13 @@ pub struct RTHMemoBuilder { fee: Amount, } +#[derive(Clone, Debug)] +pub enum CustomMemoType { + PaymentRequestId(u64), + PaymentIntentId(u64), + FlexibleMemoGenerator(FlexibleMemoGenerator), +} + /// A function that returns a flexible memopayload for outputs pub type FlexibleOutputMemoGenerator = Box< dyn Fn(FlexibleMemoOutputContext) -> Result + Sync + Send, @@ -137,9 +140,7 @@ impl Default for RTHMemoBuilder { fn default() -> Self { Self { sender_cred: Default::default(), - payment_request_id: None, - payment_intent_id: None, - flexible_memo_generator: None, + custom_memo_type: None, destination_memo_enabled: false, wrote_destination_memo: false, last_recipient: Default::default(), @@ -151,6 +152,18 @@ impl Default for RTHMemoBuilder { } } +/// An error that occurs when setting up a memo builder +/// +/// These errors are usually created setting invalid field combinations on the +/// memo builder. We have included error codes for some known useful error +/// conditions. For a custom MemoBuilder, you can try to reuse those, or use the +/// Other error code. +#[derive(Clone, Debug, Display, Eq, PartialEq)] +pub enum MemoBuilderError { + /// Invalid state change + StateChange(String), +} + impl RTHMemoBuilder { /// Set the sender credential. If no sender credential is provided, /// then authenticated sender memos cannot be produced. @@ -176,33 +189,47 @@ impl RTHMemoBuilder { } /// Set the payment request id. - pub fn set_payment_request_id(&mut self, id: u64) { - self.payment_request_id = Some(id); + pub fn set_payment_request_id(&mut self, id: u64) -> Result<(), MemoBuilderError> { + if self.custom_memo_type.is_some() { + return Err(MemoBuilderError::StateChange(format!( + "Custom Memo Type already set to {:?}", + self.custom_memo_type + ))); + } + self.custom_memo_type = Some(CustomMemoType::PaymentRequestId(id)); + Ok(()) } - /// Clear the payment request id. - pub fn clear_payment_request_id(&mut self) { - self.payment_request_id = None; + /// Clear the custom memo type. + pub fn clear_custom_memo_type(&mut self) { + self.custom_memo_type = None; } /// Set the payment intent id. - pub fn set_payment_intent_id(&mut self, id: u64) { - self.payment_intent_id = Some(id); - } - - /// Clear the payment intent id. - pub fn clear_payment_intent_id(&mut self) { - self.payment_intent_id = None; + pub fn set_payment_intent_id(&mut self, id: u64) -> Result<(), MemoBuilderError> { + if self.custom_memo_type.is_some() { + return Err(MemoBuilderError::StateChange(format!( + "Custom Memo Type already set to {:?}", + self.custom_memo_type + ))); + } + self.custom_memo_type = Some(CustomMemoType::PaymentIntentId(id)); + Ok(()) } /// Set the flexible memo generator. - pub fn set_flexible_memo_generator(&mut self, generator: FlexibleMemoGenerator) { - self.flexible_memo_generator = Some(generator); - } - - /// clear the flexible memo generator. - pub fn clear_flexible_memo_generator(&mut self) { - self.flexible_memo_generator = None; + pub fn set_flexible_memo_generator( + &mut self, + generator: FlexibleMemoGenerator, + ) -> Result<(), MemoBuilderError> { + if self.custom_memo_type.is_some() { + return Err(MemoBuilderError::StateChange(format!( + "Custom Memo Type already set to {:?}", + self.custom_memo_type + ))); + } + self.custom_memo_type = Some(CustomMemoType::FlexibleMemoGenerator(generator)); + Ok(()) } /// Enable destination memos @@ -269,53 +296,51 @@ impl MemoBuilder for RTHMemoBuilder { self.last_recipient = ShortAddressHash::from(recipient); let payload: MemoPayload = if let Some(cred) = &self.sender_cred { - if self.flexible_memo_generator.is_some() { - let tx_public_key = memo_context.tx_public_key; - let flexible_memo_context = FlexibleMemoOutputContext { - memo_context, - amount, - recipient, - builder_context: self.get_flexible_memo_builder_context(), - }; - let flexible_memo_payload = - (self - .flexible_memo_generator - .as_ref() - .unwrap() - .flexible_output_memo_generator)(flexible_memo_context)?; - let memo_data = compute_authenticated_sender_memo( - flexible_memo_payload.memo_type_bytes, - cred, - recipient.view_public_key(), - &tx_public_key.into(), - &flexible_memo_payload.memo_data, - ); - if flexible_memo_payload.memo_type_bytes[0] != 0x01 { - return Err(NewMemoError::FlexibleMemoGenerator(format!("The flexible output memo generator created a memo of the incorrect memo type: {:?}", flexible_memo_payload.memo_type_bytes))); + if let Some(custom_memo_type) = &self.custom_memo_type { + match custom_memo_type { + CustomMemoType::FlexibleMemoGenerator(flexible_memo_generator) => { + let tx_public_key = memo_context.tx_public_key; + let flexible_memo_context = FlexibleMemoOutputContext { + memo_context, + amount, + recipient, + builder_context: self.get_flexible_memo_builder_context(), + }; + let flexible_memo_payload = (flexible_memo_generator + .flexible_output_memo_generator)( + flexible_memo_context + )?; + let memo_data = compute_authenticated_sender_memo( + flexible_memo_payload.memo_type_bytes, + cred, + recipient.view_public_key(), + &tx_public_key.into(), + &flexible_memo_payload.memo_data, + ); + if flexible_memo_payload.memo_type_bytes[0] != 0x01 { + return Err(NewMemoError::FlexibleMemoGenerator(format!("The flexible output memo generator created a memo of the incorrect memo type: {:?}", flexible_memo_payload.memo_type_bytes))); + } + MemoPayload::new(flexible_memo_payload.memo_type_bytes, memo_data) + } + CustomMemoType::PaymentRequestId(payment_request_id) => { + AuthenticatedSenderWithPaymentRequestIdMemo::new( + cred, + recipient.view_public_key(), + &memo_context.tx_public_key.into(), + *payment_request_id, + ) + .into() + } + CustomMemoType::PaymentIntentId(payment_intent_id) => { + AuthenticatedSenderWithPaymentIntentIdMemo::new( + cred, + recipient.view_public_key(), + &memo_context.tx_public_key.into(), + *payment_intent_id, + ) + .into() + } } - let payload: MemoPayload = - MemoPayload::new(flexible_memo_payload.memo_type_bytes, memo_data); - return Ok(payload); - } - if self.payment_request_id.is_some() && self.payment_intent_id.is_some() { - return Err(NewMemoError::RequestAndIntentIdSet); - } - if let Some(payment_request_id) = self.payment_request_id { - AuthenticatedSenderWithPaymentRequestIdMemo::new( - cred, - recipient.view_public_key(), - &memo_context.tx_public_key.into(), - payment_request_id, - ) - .into() - } else if let Some(payment_intent_id) = self.payment_intent_id { - AuthenticatedSenderWithPaymentIntentIdMemo::new( - cred, - recipient.view_public_key(), - &memo_context.tx_public_key.into(), - payment_intent_id, - ) - .into() } else { AuthenticatedSenderMemo::new( cred, @@ -364,71 +389,73 @@ impl MemoBuilder for RTHMemoBuilder { .total_outlay .checked_add(self.fee.value) .ok_or(NewMemoError::LimitsExceeded("total_outlay"))?; - if self.flexible_memo_generator.is_some() { - let flexible_memo_context = FlexibleMemoChangeContext { - memo_context: _memo_context, - amount, - change_destination: _change_destination, - builder_context: self.get_flexible_memo_builder_context(), - }; - let flexible_memo_payload = - (self - .flexible_memo_generator - .as_ref() - .unwrap() - .flexible_change_memo_generator)(flexible_memo_context)?; - if flexible_memo_payload.memo_type_bytes[0] != 0x02 { - return Err(NewMemoError::FlexibleMemoGenerator(format!("The flexible change memo generator created a memo of the incorrect memo type: {:?}", flexible_memo_payload.memo_type_bytes))); - } - let memo_data = compute_destination_memo( - self.last_recipient.clone(), - self.fee.value, - self.num_recipients, - self.total_outlay, - flexible_memo_payload.memo_data, - ); - let payload: MemoPayload = - MemoPayload::new(flexible_memo_payload.memo_type_bytes, memo_data); - //TODO: Check whether fee is too large before setting wrote destination memo - self.wrote_destination_memo = true; - return Ok(payload); - } - - if self.payment_request_id.is_some() && self.payment_intent_id.is_some() { - return Err(NewMemoError::RequestAndIntentIdSet); - } - - if let Some(payment_request_id) = self.payment_request_id { - match DestinationWithPaymentRequestIdMemo::new( - self.last_recipient.clone(), - self.total_outlay, - self.fee.value, - payment_request_id, - ) { - Ok(mut d_memo) => { + if let Some(custom_memo_type) = &self.custom_memo_type { + match custom_memo_type { + CustomMemoType::FlexibleMemoGenerator(flexible_memo_generator) => { + let flexible_memo_context = FlexibleMemoChangeContext { + memo_context: _memo_context, + amount, + change_destination: _change_destination, + builder_context: self.get_flexible_memo_builder_context(), + }; + let flexible_memo_payload = (flexible_memo_generator + .flexible_change_memo_generator)( + flexible_memo_context + )?; + if flexible_memo_payload.memo_type_bytes[0] != 0x02 { + return Err(NewMemoError::FlexibleMemoGenerator(format!("The flexible change memo generator created a memo of the incorrect memo type: {:?}", flexible_memo_payload.memo_type_bytes))); + } + let memo_data = compute_destination_memo( + self.last_recipient.clone(), + self.fee.value, + self.num_recipients, + self.total_outlay, + flexible_memo_payload.memo_data, + ); + let payload: MemoPayload = + MemoPayload::new(flexible_memo_payload.memo_type_bytes, memo_data); + //TODO: Check whether fee is too large before setting wrote destination memo self.wrote_destination_memo = true; - d_memo.set_num_recipients(self.num_recipients); - Ok(d_memo.into()) + Ok(payload) } - Err(err) => match err { - DestinationMemoError::FeeTooLarge => Err(NewMemoError::LimitsExceeded("fee")), - }, - } - } else if let Some(payment_intent_id) = self.payment_intent_id { - match DestinationWithPaymentIntentIdMemo::new( - self.last_recipient.clone(), - self.total_outlay, - self.fee.value, - payment_intent_id, - ) { - Ok(mut d_memo) => { - self.wrote_destination_memo = true; - d_memo.set_num_recipients(self.num_recipients); - Ok(d_memo.into()) + CustomMemoType::PaymentRequestId(payment_request_id) => { + match DestinationWithPaymentRequestIdMemo::new( + self.last_recipient.clone(), + self.total_outlay, + self.fee.value, + *payment_request_id, + ) { + Ok(mut d_memo) => { + self.wrote_destination_memo = true; + d_memo.set_num_recipients(self.num_recipients); + Ok(d_memo.into()) + } + Err(err) => match err { + DestinationMemoError::FeeTooLarge => { + Err(NewMemoError::LimitsExceeded("fee")) + } + }, + } + } + CustomMemoType::PaymentIntentId(payment_intent_id) => { + match DestinationWithPaymentIntentIdMemo::new( + self.last_recipient.clone(), + self.total_outlay, + self.fee.value, + *payment_intent_id, + ) { + Ok(mut d_memo) => { + self.wrote_destination_memo = true; + d_memo.set_num_recipients(self.num_recipients); + Ok(d_memo.into()) + } + Err(err) => match err { + DestinationMemoError::FeeTooLarge => { + Err(NewMemoError::LimitsExceeded("fee")) + } + }, + } } - Err(err) => match err { - DestinationMemoError::FeeTooLarge => Err(NewMemoError::LimitsExceeded("fee")), - }, } } else { match DestinationMemo::new( @@ -630,7 +657,9 @@ mod tests { let mut rng: StdRng = SeedableRng::from_seed([0u8; 32]); let (sender, mut builder) = build_test_memo_builder(&mut rng); // Add a flexible memo generator - builder.set_flexible_memo_generator(get_valid_flexible_memo_generator()); + builder + .set_flexible_memo_generator(get_valid_flexible_memo_generator()) + .expect("No other custom memo type should be set"); // Build the memo payload let memo_test_context = build_rth_memos(sender, builder, funding_amount, change_amount, fee); @@ -679,7 +708,9 @@ mod tests { let mut rng: StdRng = SeedableRng::from_seed([0u8; 32]); let (sender, mut builder) = build_test_memo_builder(&mut rng); // Add a flexible memo generator - builder.set_flexible_memo_generator(get_flexible_memo_generator_returning_invalid_types()); + builder + .set_flexible_memo_generator(get_flexible_memo_generator_returning_invalid_types()) + .expect("No other custom memo types should be set"); // Build the memo payload let memo_test_context = build_rth_memos(sender, builder, funding_amount, change_amount, fee); diff --git a/transaction/builder/src/transaction_builder.rs b/transaction/builder/src/transaction_builder.rs index a066cdc90b..c22a2fa526 100644 --- a/transaction/builder/src/transaction_builder.rs +++ b/transaction/builder/src/transaction_builder.rs @@ -1844,7 +1844,9 @@ pub mod transaction_builder_tests { let mut memo_builder = RTHMemoBuilder::default(); memo_builder.set_sender_credential(SenderMemoCredential::from(&sender)); memo_builder.enable_destination_memo(); - memo_builder.set_payment_request_id(42); + memo_builder + .set_payment_request_id(42) + .expect("No other memo types should be set"); let mut transaction_builder = TransactionBuilder::new( block_version, @@ -2009,7 +2011,9 @@ pub mod transaction_builder_tests { { let mut memo_builder = RTHMemoBuilder::default(); memo_builder.set_sender_credential(SenderMemoCredential::from(&sender)); - memo_builder.set_payment_request_id(47); + memo_builder + .set_payment_request_id(47) + .expect("No other memo types should be set"); let mut transaction_builder = TransactionBuilder::new( block_version, @@ -2161,7 +2165,9 @@ pub mod transaction_builder_tests { { let mut memo_builder = RTHMemoBuilder::default(); memo_builder.enable_destination_memo(); - memo_builder.set_payment_request_id(47); + memo_builder + .set_payment_request_id(47) + .expect("No other memo types should be set"); let mut transaction_builder = TransactionBuilder::new( block_version, @@ -2309,7 +2315,9 @@ pub mod transaction_builder_tests { let mut memo_builder = RTHMemoBuilder::default(); memo_builder.set_sender_credential(SenderMemoCredential::from(&sender)); memo_builder.enable_destination_memo(); - memo_builder.set_payment_intent_id(4855282172840142080); + memo_builder + .set_payment_intent_id(4855282172840142080) + .expect("No other memo types should be set"); let mut transaction_builder = TransactionBuilder::new( block_version,