From 414de6dbe285f18f7946c8e3b91ccef340240388 Mon Sep 17 00:00:00 2001 From: Gianmarco Fraccaroli <> Date: Thu, 25 Jan 2024 12:07:47 +0100 Subject: [PATCH 1/2] refactor multitoken vp with checked arithmetic --- .../namada/src/ledger/native_vp/multitoken.rs | 74 +++++++++++++++---- 1 file changed, 61 insertions(+), 13 deletions(-) diff --git a/crates/namada/src/ledger/native_vp/multitoken.rs b/crates/namada/src/ledger/native_vp/multitoken.rs index 5537ca7541..ed07c93256 100644 --- a/crates/namada/src/ledger/native_vp/multitoken.rs +++ b/crates/namada/src/ledger/native_vp/multitoken.rs @@ -1,7 +1,9 @@ //! Native VP for multitokens use std::collections::{BTreeSet, HashMap}; +use std::ops::Neg; +use namada_core::types::uint::I256; use namada_tx::Tx; use namada_vp_env::VpEnv; use thiserror::Error; @@ -21,11 +23,20 @@ use crate::vm::WasmCacheAccess; pub enum Error { #[error("Native VP error: {0}")] NativeVpError(#[from] native_vp::Error), + #[error("Couldn't read token amount")] + InvalidTokenAmountKey, + #[error("Amount subtraction underflowed")] + AmountUnderflow, } /// Multitoken functions result pub type Result = std::result::Result; +enum ReadType { + Pre, + Post, +} + /// Multitoken VP pub struct MultitokenVp<'a, DB, H, CA> where @@ -55,26 +66,31 @@ where let mut mints = HashMap::new(); for key in keys_changed { if let Some([token, _]) = is_any_token_balance_key(key) { - let pre: Amount = self.ctx.read_pre(key)?.unwrap_or_default(); - let post: Amount = self.ctx.read_post(key)?.unwrap_or_default(); - let diff = post.change() - pre.change(); + let pre_amount = self.read_amount(key, ReadType::Pre)?; + let post_amount = self.read_amount(key, ReadType::Post)?; + println!("{}, {}", pre_amount, post_amount); + let diff_amount = + self.compute_post_pre_diff(pre_amount, post_amount)?; + match changes.get_mut(token) { - Some(change) => *change += diff, - None => _ = changes.insert(token, diff), + Some(change) => *change += diff_amount, + None => _ = changes.insert(token, diff_amount), } } else if let Some(token) = is_any_minted_balance_key(key) { - let pre: Amount = self.ctx.read_pre(key)?.unwrap_or_default(); - let post: Amount = self.ctx.read_post(key)?.unwrap_or_default(); - let diff = post.change() - pre.change(); - match mints.get_mut(token) { - Some(mint) => *mint += diff, - None => _ = mints.insert(token, diff), - } - // Check if the minter is set if !self.is_valid_minter(token, verifiers)? { return Ok(false); } + + let pre_amount = self.read_amount(key, ReadType::Pre)?; + let post_amount = self.read_amount(key, ReadType::Post)?; + let diff_amount = + self.compute_post_pre_diff(pre_amount, post_amount)?; + + match mints.get_mut(token) { + Some(mint) => *mint += diff_amount, + None => _ = mints.insert(token, diff_amount), + } } else if let Some(token) = is_any_minter_key(key) { if !self.is_valid_minter(token, verifiers)? { return Ok(false); @@ -133,6 +149,38 @@ where } } } + + fn read_amount(&self, key: &Key, read_type: ReadType) -> Result { + let result = match read_type { + ReadType::Pre => self.ctx.read_pre::(key), + ReadType::Post => self.ctx.read_post::(key), + }; + + match result { + Ok(Some(amount)) => Ok(amount), + Ok(None) => Ok(Amount::zero()), + _ => Err(Error::InvalidTokenAmountKey), + } + } + + // this function computes the difference between post and pre amount + fn compute_post_pre_diff( + &self, + amount_pre: Amount, + amount_post: Amount, + ) -> Result { + if amount_pre > amount_post { + match amount_pre.checked_sub(amount_post) { + Some(diff) => Ok(diff.change().neg()), + None => Err(Error::AmountUnderflow), + } + } else { + match amount_post.checked_sub(amount_pre) { + Some(diff) => Ok(diff.change()), + None => Err(Error::AmountUnderflow), + } + } + } } #[cfg(test)] From d1d166077bd8db87f395c85434f4fb069ce41c67 Mon Sep 17 00:00:00 2001 From: Gianmarco Fraccaroli <> Date: Thu, 25 Jan 2024 13:40:51 +0100 Subject: [PATCH 2/2] refactor multitoken vp with checked arithmetic --- Cargo.lock | 1 + crates/namada/Cargo.toml | 1 + .../namada/src/ledger/native_vp/multitoken.rs | 18 +++++++++++++----- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 71515d9ae8..99df830681 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4162,6 +4162,7 @@ dependencies = [ "namada_tx_env", "namada_vote_ext", "namada_vp_env", + "num-traits 0.2.17", "num256", "orion", "owo-colors", diff --git a/crates/namada/Cargo.toml b/crates/namada/Cargo.toml index 0180115988..f818917863 100644 --- a/crates/namada/Cargo.toml +++ b/crates/namada/Cargo.toml @@ -144,6 +144,7 @@ wasmer-vm = { git = "https://github.com/heliaxdev/wasmer", rev = "255054f7f58b7b wat = "=1.0.71" wasmparser.workspace = true zeroize.workspace = true +num-traits.workspace = true [target.'cfg(not(target_family = "wasm"))'.dependencies] tokio = { workspace = true, features = ["full"] } diff --git a/crates/namada/src/ledger/native_vp/multitoken.rs b/crates/namada/src/ledger/native_vp/multitoken.rs index ed07c93256..6eb87d319a 100644 --- a/crates/namada/src/ledger/native_vp/multitoken.rs +++ b/crates/namada/src/ledger/native_vp/multitoken.rs @@ -6,6 +6,7 @@ use std::ops::Neg; use namada_core::types::uint::I256; use namada_tx::Tx; use namada_vp_env::VpEnv; +use num_traits::ops::checked::CheckedAdd; use thiserror::Error; use crate::ledger::native_vp::{self, Ctx, NativeVp}; @@ -27,6 +28,8 @@ pub enum Error { InvalidTokenAmountKey, #[error("Amount subtraction underflowed")] AmountUnderflow, + #[error("Amount subtraction overflowed")] + AmountOverflowed, } /// Multitoken functions result @@ -62,18 +65,20 @@ where keys_changed: &BTreeSet, verifiers: &BTreeSet
, ) -> Result { - let mut changes = HashMap::new(); - let mut mints = HashMap::new(); + let mut changes: HashMap<&Address, I256> = HashMap::new(); + let mut mints: HashMap<&Address, I256> = HashMap::new(); for key in keys_changed { if let Some([token, _]) = is_any_token_balance_key(key) { let pre_amount = self.read_amount(key, ReadType::Pre)?; let post_amount = self.read_amount(key, ReadType::Post)?; - println!("{}, {}", pre_amount, post_amount); let diff_amount = self.compute_post_pre_diff(pre_amount, post_amount)?; match changes.get_mut(token) { - Some(change) => *change += diff_amount, + Some(change) => match change.checked_add(&diff_amount) { + Some(new_change) => *change = new_change, + None => return Err(Error::AmountOverflowed), + }, None => _ = changes.insert(token, diff_amount), } } else if let Some(token) = is_any_minted_balance_key(key) { @@ -88,7 +93,10 @@ where self.compute_post_pre_diff(pre_amount, post_amount)?; match mints.get_mut(token) { - Some(mint) => *mint += diff_amount, + Some(mint) => match mint.checked_add(&diff_amount) { + Some(new_mint) => *mint = new_mint, + None => return Err(Error::AmountOverflowed), + }, None => _ = mints.insert(token, diff_amount), } } else if let Some(token) = is_any_minter_key(key) {