Whitelisted gas in wasm

klatos063 · Aug 21, 2023 · 2ca6fbd · 2ca6fbd
1 parent 72e1c62
commit 2ca6fbd
Show file tree

Hide file tree

Showing 18 changed files with 89 additions and 34 deletions.
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
@@ -10,10 +10,14 @@ use proc_macro::TokenStream;
 use proc_macro2::{Span as Span2, TokenStream as TokenStream2};
 use quote::{quote, ToTokens};
 use syn::punctuated::Punctuated;
-use syn::{parse_macro_input, ItemFn, ItemStruct};
+use syn::{parse_macro_input, ExprAssign, FnArg, ItemFn, ItemStruct, Pat};
 
 /// Generate WASM binding for a transaction main entrypoint function.
 ///
+/// It expects an attribute in the form: `gas = u64`, so that a call to the gas
+/// meter can be injected as the first instruction of the transaction to account
+/// for the whitelisted gas amount.
+///
 /// This macro expects a function with signature:
 ///
 /// ```compiler_fail
@@ -23,15 +27,38 @@ use syn::{parse_macro_input, ItemFn, ItemStruct};
 /// ) -> TxResult
 /// ```
 #[proc_macro_attribute]
-pub fn transaction(_attr: TokenStream, input: TokenStream) -> TokenStream {
+pub fn transaction(attr: TokenStream, input: TokenStream) -> TokenStream {
     let ast = parse_macro_input!(input as ItemFn);
-    let ident = &ast.sig.ident;
+    let ItemFn {
+        attrs,
+        vis,
+        sig,
+        block,
+    } = ast;
+    let stmts = &block.stmts;
+    let ident = &sig.ident;
+    let attr_ast = parse_macro_input!(attr as ExprAssign);
+    let gas = attr_ast.right;
+    let ctx = match sig.inputs.first() {
+        Some(FnArg::Typed(pat_type)) => {
+            if let Pat::Ident(pat_ident) = pat_type.pat.as_ref() {
+                &pat_ident.ident
+            } else {
+                panic!("Unexpected token, expected ctx ident")
+            }
+        }
+        _ => panic!("Unexpected token, expected ctx ident"),
+    };
     let gen = quote! {
         // Use `wee_alloc` as the global allocator.
         #[global_allocator]
         static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
 
-        #ast
+        #(#attrs)* #vis #sig {
+            // Consume the whitelisted gas
+            #ctx.charge_gas(#gas)?;
+            #(#stmts)*
+        }
 
         // The module entrypoint callable by wasm runtime
         #[no_mangle]
@@ -63,6 +90,10 @@ pub fn transaction(_attr: TokenStream, input: TokenStream) -> TokenStream {
 
 /// Generate WASM binding for validity predicate main entrypoint function.
 ///
+/// It expects an attribute in the form: `gas = u64`, so that a call to the gas
+/// meter can be injected as the first instruction of the validity predicate to
+/// account for the whitelisted gas amount.
+///
 /// This macro expects a function with signature:
 ///
 /// ```compiler_fail
@@ -76,17 +107,40 @@ pub fn transaction(_attr: TokenStream, input: TokenStream) -> TokenStream {
 /// ```
 #[proc_macro_attribute]
 pub fn validity_predicate(
-    _attr: TokenStream,
+    attr: TokenStream,
     input: TokenStream,
 ) -> TokenStream {
     let ast = parse_macro_input!(input as ItemFn);
-    let ident = &ast.sig.ident;
+    let ItemFn {
+        attrs,
+        vis,
+        sig,
+        block,
+    } = ast;
+    let stmts = &block.stmts;
+    let ident = &sig.ident;
+    let attr_ast = parse_macro_input!(attr as ExprAssign);
+    let gas = attr_ast.right;
+    let ctx = match sig.inputs.first() {
+        Some(FnArg::Typed(pat_type)) => {
+            if let Pat::Ident(pat_ident) = pat_type.pat.as_ref() {
+                &pat_ident.ident
+            } else {
+                panic!("Unexpected token, expected ctx ident")
+            }
+        }
+        _ => panic!("Unexpected token, expected ctx ident"),
+    };
     let gen = quote! {
         // Use `wee_alloc` as the global allocator.
         #[global_allocator]
         static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
 
-        #ast
+        #(#attrs)* #vis #sig {
+            // Consume the whitelisted gas
+            #ctx.charge_gas(#gas)?;
+            #(#stmts)*
+        }
 
         // The module entrypoint callable by wasm runtime
         #[no_mangle]

diff --git a/wasm/tx_template/src/lib.rs b/wasm/tx_template/src/lib.rs
@@ -1,6 +1,6 @@
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 1000)]
 fn apply_tx(_ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     log_string(format!("apply_tx called with data: {:#?}", tx_data));
     Ok(())

diff --git a/wasm/vp_template/src/lib.rs b/wasm/vp_template/src/lib.rs
@@ -1,6 +1,6 @@
 use namada_vp_prelude::*;
 
-#[validity_predicate]
+#[validity_predicate(gas = 1000)]
 fn validate_tx(
     ctx: &Ctx,
     tx_data: Tx,

diff --git a/wasm/wasm_source/src/tx_bond.rs b/wasm/wasm_source/src/tx_bond.rs
@@ -2,7 +2,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 160000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_bridge_pool.rs b/wasm/wasm_source/src/tx_bridge_pool.rs
@@ -5,7 +5,7 @@ use eth_bridge::storage::{bridge_pool, native_erc20_key, wrapped_erc20s};
 use eth_bridge_pool::{GasFee, PendingTransfer, TransferToEthereum};
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 100000)]
 fn apply_tx(ctx: &mut Ctx, signed: Tx) -> TxResult {
     let data = signed.data().ok_or_err_msg("Missing data")?;
     let transfer = PendingTransfer::try_from_slice(&data[..])

diff --git a/wasm/wasm_source/src/tx_change_validator_commission.rs b/wasm/wasm_source/src/tx_change_validator_commission.rs
@@ -3,7 +3,7 @@
 use namada_tx_prelude::transaction::pos::CommissionChange;
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 220000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_ibc.rs b/wasm/wasm_source/src/tx_ibc.rs
@@ -5,7 +5,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 1240000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_init_account.rs b/wasm/wasm_source/src/tx_init_account.rs
@@ -3,7 +3,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 230000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_init_proposal.rs b/wasm/wasm_source/src/tx_init_proposal.rs
@@ -2,7 +2,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 40000)]
 fn apply_tx(ctx: &mut Ctx, tx: Tx) -> TxResult {
     let data = tx.data().ok_or_err_msg("Missing data")?;
     let tx_data =

diff --git a/wasm/wasm_source/src/tx_init_validator.rs b/wasm/wasm_source/src/tx_init_validator.rs
@@ -4,7 +4,7 @@
 use namada_tx_prelude::transaction::pos::InitValidator;
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 730000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_reveal_pk.rs b/wasm/wasm_source/src/tx_reveal_pk.rs
@@ -6,7 +6,7 @@
 use namada_tx_prelude::key::common;
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 170000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_transfer.rs b/wasm/wasm_source/src/tx_transfer.rs
@@ -4,7 +4,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 110000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_unbond.rs b/wasm/wasm_source/src/tx_unbond.rs
@@ -3,7 +3,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 430000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_unjail_validator.rs b/wasm/wasm_source/src/tx_unjail_validator.rs
@@ -3,7 +3,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 340000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_update_account.rs b/wasm/wasm_source/src/tx_update_account.rs
@@ -5,7 +5,7 @@
 use namada_tx_prelude::key::pks_handle;
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 140000)]
 fn apply_tx(ctx: &mut Ctx, tx: Tx) -> TxResult {
     let signed = tx;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_vote_proposal.rs b/wasm/wasm_source/src/tx_vote_proposal.rs
@@ -2,7 +2,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 120000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm/wasm_source/src/tx_withdraw.rs b/wasm/wasm_source/src/tx_withdraw.rs
@@ -3,7 +3,7 @@
 
 use namada_tx_prelude::*;
 
-#[transaction]
+#[transaction(gas = 260000)]
 fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
     let signed = tx_data;
     let data = signed.data().ok_or_err_msg("Missing data")?;

diff --git a/wasm_for_tests/wasm_source/src/lib.rs b/wasm_for_tests/wasm_source/src/lib.rs
@@ -3,7 +3,7 @@
 pub mod main {
     use namada_tx_prelude::*;
 
-    #[transaction]
+    #[transaction(gas = 1000)]
     fn apply_tx(_ctx: &mut Ctx, _tx_data: Tx) -> TxResult {
         Ok(())
     }
@@ -14,7 +14,7 @@ pub mod main {
 pub mod main {
     use namada_tx_prelude::*;
 
-    #[transaction]
+    #[transaction(gas = 1000)]
     fn apply_tx(_ctx: &mut Ctx, tx_data: Tx) -> TxResult {
         let len = usize::try_from_slice(&tx_data.data().as_ref().unwrap()[..])
             .unwrap();
@@ -31,7 +31,7 @@ pub mod main {
 pub mod main {
     use namada_tx_prelude::*;
 
-    #[transaction]
+    #[transaction(gas = 1000)]
     fn apply_tx(ctx: &mut Ctx, _tx_data: Tx) -> TxResult {
         // governance
         let target_key = gov_storage::keys::get_min_proposal_grace_epoch_key();
@@ -49,7 +49,7 @@ pub mod main {
 pub mod main {
     use namada_tx_prelude::*;
 
-    #[transaction]
+    #[transaction(gas = 1000)]
     fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
         // Allocates a memory of size given from the `tx_data (usize)`
         let key =
@@ -67,7 +67,8 @@ pub mod main {
     use borsh::BorshDeserialize;
     use namada_test_utils::tx_data::TxWriteData;
     use namada_tx_prelude::{
-        log_string, transaction, Ctx, StorageRead, StorageWrite, Tx, TxResult,
+        log_string, transaction, Ctx, StorageRead, StorageWrite, Tx, TxEnv,
+        TxResult,
     };
 
     const TX_NAME: &str = "tx_write";
@@ -86,7 +87,7 @@ pub mod main {
         panic!()
     }
 
-    #[transaction]
+    #[transaction(gas = 1000)]
     fn apply_tx(ctx: &mut Ctx, tx_data: Tx) -> TxResult {
         let signed = tx_data;
         let data = match signed.data() {
@@ -132,7 +133,7 @@ pub mod main {
 pub mod main {
     use namada_vp_prelude::*;
 
-    #[validity_predicate]
+    #[validity_predicate(gas = 1000)]
     fn validate_tx(
         _ctx: &Ctx,
         _tx_data: Tx,
@@ -149,7 +150,7 @@ pub mod main {
 pub mod main {
     use namada_vp_prelude::*;
 
-    #[validity_predicate]
+    #[validity_predicate(gas = 1000)]
     fn validate_tx(
         _ctx: &Ctx,
         _tx_data: Tx,
@@ -167,7 +168,7 @@ pub mod main {
 pub mod main {
     use namada_vp_prelude::*;
 
-    #[validity_predicate]
+    #[validity_predicate(gas = 1000)]
     fn validate_tx(
         ctx: &Ctx,
         tx_data: Tx,
@@ -192,7 +193,7 @@ pub mod main {
 pub mod main {
     use namada_vp_prelude::*;
 
-    #[validity_predicate]
+    #[validity_predicate(gas = 1000)]
     fn validate_tx(
         _ctx: &Ctx,
         tx_data: Tx,
@@ -216,7 +217,7 @@ pub mod main {
 pub mod main {
     use namada_vp_prelude::*;
 
-    #[validity_predicate]
+    #[validity_predicate(gas = 1000)]
     fn validate_tx(
         ctx: &Ctx,
         tx_data: Tx,