From a6c8c3ee1461a53ed78574ff19955a229f87570f Mon Sep 17 00:00:00 2001 From: Milosz Muszynski Date: Thu, 18 Jan 2024 11:06:49 +0100 Subject: [PATCH] more protection of logic calls, added num_notes --- contracts/transfer-data/src/lib.rs | 8 ++++ contracts/transfer-data/src/state.rs | 5 +++ contracts/transfer-data/src/tree.rs | 4 ++ contracts/transfer-proxy/src/lib.rs | 5 +++ contracts/transfer-proxy/src/state.rs | 6 +++ contracts/transfer/src/lib.rs | 64 +++++++++++++++++++++------ contracts/transfer/src/state.rs | 3 +- 7 files changed, 80 insertions(+), 15 deletions(-) diff --git a/contracts/transfer-data/src/lib.rs b/contracts/transfer-data/src/lib.rs index 7d8c879957..a2eae98cc6 100644 --- a/contracts/transfer-data/src/lib.rs +++ b/contracts/transfer-data/src/lib.rs @@ -33,6 +33,14 @@ unsafe fn root(arg_len: u32) -> u32 { }) } +#[no_mangle] +unsafe fn num_notes(arg_len: u32) -> u32 { + rusk_abi::wrap_call(arg_len, |_: ()| { + assert_transfer_caller(); + STATE.num_notes() + }) +} + #[no_mangle] unsafe fn module_balance(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |module| { diff --git a/contracts/transfer-data/src/state.rs b/contracts/transfer-data/src/state.rs index fd8e58576e..9d497bbd40 100644 --- a/contracts/transfer-data/src/state.rs +++ b/contracts/transfer-data/src/state.rs @@ -83,6 +83,11 @@ impl TransferState { self.tree.root() } + /// Get the count of the notes in the tree. + pub fn num_notes(&self) -> u64 { + self.tree.leaves_len() + } + /// Get the opening pub fn opening( &self, diff --git a/contracts/transfer-data/src/tree.rs b/contracts/transfer-data/src/tree.rs index ad9c8c5858..9bc3d974e7 100644 --- a/contracts/transfer-data/src/tree.rs +++ b/contracts/transfer-data/src/tree.rs @@ -93,4 +93,8 @@ impl Tree { ) -> Option> { self.tree.opening(pos) } + + pub fn leaves_len(&self) -> u64 { + self.leaves.len() as u64 + } } diff --git a/contracts/transfer-proxy/src/lib.rs b/contracts/transfer-proxy/src/lib.rs index 0ae99663f8..79807ad8f1 100644 --- a/contracts/transfer-proxy/src/lib.rs +++ b/contracts/transfer-proxy/src/lib.rs @@ -84,6 +84,11 @@ unsafe fn root(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |_: ()| STATE.root()) } +#[no_mangle] +unsafe fn num_notes(arg_len: u32) -> u32 { + rusk_abi::wrap_call(arg_len, |_: ()| STATE.num_notes()) +} + #[no_mangle] unsafe fn module_balance(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |module| STATE.balance(&module)) diff --git a/contracts/transfer-proxy/src/state.rs b/contracts/transfer-proxy/src/state.rs index 212560444c..822ce08ff5 100644 --- a/contracts/transfer-proxy/src/state.rs +++ b/contracts/transfer-proxy/src/state.rs @@ -188,6 +188,12 @@ impl TransferProxy { .expect("root query should succeed") } + /// Get the count of the notes in the tree. + pub fn num_notes(&self) -> u64 { + rusk_abi::call::<(), u64>(self.target, "num_notes", &()) + .expect("num_notes query should succeed") + } + /// Get the opening pub fn opening( &self, diff --git a/contracts/transfer/src/lib.rs b/contracts/transfer/src/lib.rs index 96be3e8d9e..74b750e151 100644 --- a/contracts/transfer/src/lib.rs +++ b/contracts/transfer/src/lib.rs @@ -28,17 +28,24 @@ static mut STATE: TransferOps = TransferOps {}; #[no_mangle] unsafe fn mint(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |arg| STATE.mint(arg)) + rusk_abi::wrap_call(arg_len, |arg| { + assert_transfer_caller(); + STATE.mint(arg) + }) } #[no_mangle] unsafe fn stct(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |arg| STATE.send_to_contract_transparent(arg)) + rusk_abi::wrap_call(arg_len, |arg| { + assert_transfer_caller(); + STATE.send_to_contract_transparent(arg) + }) } #[no_mangle] unsafe fn wfct(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |(arg, from_address)| { + assert_transfer_caller(); STATE.withdraw_from_contract_transparent(arg, from_address) }) } @@ -46,18 +53,23 @@ unsafe fn wfct(arg_len: u32) -> u32 { #[no_mangle] unsafe fn wfct_raw(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |(arg, from_address)| { + assert_transfer_caller(); STATE.withdraw_from_contract_transparent_raw(arg, from_address) }) } #[no_mangle] unsafe fn stco(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |arg| STATE.send_to_contract_obfuscated(arg)) + rusk_abi::wrap_call(arg_len, |arg| { + assert_transfer_caller(); + STATE.send_to_contract_obfuscated(arg) + }) } #[no_mangle] unsafe fn wfco(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |(arg, from_address)| { + assert_transfer_caller(); STATE.withdraw_from_contract_obfuscated(arg, from_address) }) } @@ -65,6 +77,7 @@ unsafe fn wfco(arg_len: u32) -> u32 { #[no_mangle] unsafe fn wfco_raw(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |(arg, from_address)| { + assert_transfer_caller(); STATE.withdraw_from_contract_obfuscated_raw(arg, from_address) }) } @@ -72,6 +85,7 @@ unsafe fn wfco_raw(arg_len: u32) -> u32 { #[no_mangle] unsafe fn wfctc(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |(arg, from_address)| { + assert_transfer_caller(); STATE.withdraw_from_contract_transparent_to_contract(arg, from_address) }) } @@ -80,46 +94,68 @@ unsafe fn wfctc(arg_len: u32) -> u32 { #[no_mangle] unsafe fn root(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |_: ()| STATE.root()) + rusk_abi::wrap_call(arg_len, |_: ()| { + assert_transfer_caller(); + STATE.root() + }) +} + +#[no_mangle] +unsafe fn num_notes(arg_len: u32) -> u32 { + rusk_abi::wrap_call(arg_len, |_: ()| { + assert_transfer_caller(); + STATE.num_notes() + }) } #[no_mangle] unsafe fn module_balance(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |module| STATE.balance(&module)) + rusk_abi::wrap_call(arg_len, |module| { + assert_transfer_caller(); + STATE.balance(&module) + }) } #[no_mangle] unsafe fn message(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |(module, pk)| STATE.message(&module, &pk)) + rusk_abi::wrap_call(arg_len, |(module, pk)| { + assert_transfer_caller(); + STATE.message(&module, &pk) + }) } #[no_mangle] unsafe fn opening(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |pos| STATE.opening(pos)) + rusk_abi::wrap_call(arg_len, |pos| { + assert_transfer_caller(); + STATE.opening(pos) + }) } #[no_mangle] unsafe fn existing_nullifiers(arg_len: u32) -> u32 { rusk_abi::wrap_call(arg_len, |nullifiers| { + assert_transfer_caller(); STATE.existing_nullifiers(&nullifiers) }) } -#[no_mangle] -unsafe fn num_notes(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |_: ()| STATE.num_notes()) -} - // "Feeder" queries #[no_mangle] unsafe fn leaves_from_height(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |height| STATE.leaves_from_height(height)) + rusk_abi::wrap_call(arg_len, |height| { + assert_transfer_caller(); + STATE.leaves_from_height(height) + }) } #[no_mangle] unsafe fn leaves_from_pos(arg_len: u32) -> u32 { - rusk_abi::wrap_call(arg_len, |pos| STATE.leaves_from_pos(pos)) + rusk_abi::wrap_call(arg_len, |pos| { + assert_transfer_caller(); + STATE.leaves_from_pos(pos) + }) } // "Management" transactions diff --git a/contracts/transfer/src/state.rs b/contracts/transfer/src/state.rs index a0ba40b943..56741cc19a 100644 --- a/contracts/transfer/src/state.rs +++ b/contracts/transfer/src/state.rs @@ -474,7 +474,8 @@ impl TransferOps { /// Get the count of the notes in the tree. pub fn num_notes(&self) -> u64 { - self.tree.leaves_len() + rusk_abi::call::<(), u64>(TRANSFER_DATA_CONTRACT, "num_notes", &()) + .expect("num_notes query should succeed") } /// Get the opening