Skip to content

Commit

Permalink
Refactor to make incompatible custom memo types into an enum
Browse files Browse the repository at this point in the history
  • Loading branch information
@wjuan-mob
wjuan-mob committed Mar 22, 2023
1 parent d583f33 commit 8a23aba
Show file tree
Hide file tree
Showing 2 changed files with 183 additions and 144 deletions.
311 changes: 171 additions & 140 deletions transaction/builder/src/memo_builder/rth_memo_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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<SenderMemoCredential>,
// The payment request id, if any
payment_request_id: Option<u64>,
// The payment intent id, if any
payment_intent_id: Option<u64>,
// The memo generator function, if any
flexible_memo_generator: Option<FlexibleMemoGenerator>,
// Different options for the custom memo data
custom_memo_type: Option<CustomMemoType>,
// Whether destination memos are enabled.
destination_memo_enabled: bool,
// Tracks if we already wrote a destination memo, for error reporting
Expand All @@ -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<FlexibleMemoPayload, NewMemoError> + Sync + Send,
Expand Down Expand Up @@ -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(),
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 8a23aba

Please sign in to comment.