From 320e3e96ff1673be75078fe81a0d28cdc6e2497c Mon Sep 17 00:00:00 2001
From: Tiago Carvalho <sugoiuguu@tfwno.gf>
Date: Wed, 8 Nov 2023 09:48:37 +0100
Subject: [PATCH] New unit tests

---
 apps/src/lib/node/ledger/shell/mod.rs         | 42 ++++++++++++++-
 .../lib/node/ledger/shell/process_proposal.rs | 52 +++++++++++++++++++
 2 files changed, 93 insertions(+), 1 deletion(-)

diff --git a/apps/src/lib/node/ledger/shell/mod.rs b/apps/src/lib/node/ledger/shell/mod.rs
index 8b93b84226..2bccb4cba5 100644
--- a/apps/src/lib/node/ledger/shell/mod.rs
+++ b/apps/src/lib/node/ledger/shell/mod.rs
@@ -1072,7 +1072,10 @@ where
         const VALID_MSG: &str = "Mempool validation passed";
         const INVALID_MSG: &str = "Mempool validation failed";
 
-        // Tx size check
+        // check tx bytes
+        //
+        // NB: always keep this as the first tx check,
+        // as it is a pretty cheap one
         if !validate_tx_bytes(&self.wl_storage, tx_bytes.len())
             .expect("Failed to get max tx bytes param from storage")
         {
@@ -2900,4 +2903,41 @@ mod shell_tests {
         );
         assert_eq!(result.code, u32::from(ErrorCodes::FeeError));
     }
+
+    /// Test max tx bytes parameter in CheckTx
+    #[test]
+    fn test_max_tx_bytes_check_tx() {
+        let (shell, _recv, _, _) = test_utils::setup();
+
+        let max_tx_bytes: u32 = {
+            let key = parameters::storage::get_max_tx_bytes_key();
+            shell
+                .wl_storage
+                .read(&key)
+                .expect("Failed to read from storage")
+                .expect("Max tx bytes should have been written to storage")
+        };
+        let new_tx = |size: u32| {
+            let keypair = super::test_utils::gen_keypair();
+            let mut tx = Tx::new(shell.chain_id.clone(), None);
+            tx.add_code("wasm_code".as_bytes().to_owned())
+                .add_data(vec![0; size as usize])
+                .sign_wrapper(keypair);
+            tx
+        };
+
+        // test a small tx
+        let result = shell.mempool_validate(
+            new_tx(50).to_bytes().as_ref(),
+            MempoolTxType::NewTransaction,
+        );
+        assert!(result.code != u32::from(ErrorCodes::TooLarge));
+
+        // max tx bytes + 1, on the other hand, is not
+        let result = shell.mempool_validate(
+            new_tx(max_tx_bytes + 1).to_bytes().as_ref(),
+            MempoolTxType::NewTransaction,
+        );
+        assert_eq!(result.code, u32::from(ErrorCodes::TooLarge));
+    }
 }
diff --git a/apps/src/lib/node/ledger/shell/process_proposal.rs b/apps/src/lib/node/ledger/shell/process_proposal.rs
index 19469d10fc..5b363bb124 100644
--- a/apps/src/lib/node/ledger/shell/process_proposal.rs
+++ b/apps/src/lib/node/ledger/shell/process_proposal.rs
@@ -459,6 +459,9 @@ where
         CA: 'static + WasmCacheAccess + Sync,
     {
         // check tx bytes
+        //
+        // NB: always keep this as the first tx check,
+        // as it is a pretty cheap one
         if !validate_tx_bytes(&self.wl_storage, tx_bytes.len())
             .expect("Failed to get max tx bytes param from storage")
         {
@@ -2760,4 +2763,53 @@ mod test_process_proposal {
             );
         }
     }
+
+    /// Test max tx bytes parameter in ProcessProposal
+    #[test]
+    fn test_max_tx_bytes_process_proposal() {
+        use namada::ledger::parameters::storage::get_max_tx_bytes_key;
+        let (shell, _recv, _, _) = test_utils::setup_at_height(3u64);
+
+        let max_tx_bytes: u32 = {
+            let key = get_max_tx_bytes_key();
+            shell
+                .wl_storage
+                .read(&key)
+                .expect("Failed to read from storage")
+                .expect("Max tx bytes should have been written to storage")
+        };
+        let new_tx = |size: u32| {
+            let keypair = super::test_utils::gen_keypair();
+            let mut tx = Tx::new(shell.chain_id.clone(), None);
+            tx.add_code("wasm_code".as_bytes().to_owned())
+                .add_data(vec![0; size as usize])
+                .sign_wrapper(keypair);
+            tx
+        };
+
+        let request = ProcessProposal {
+            txs: vec![new_tx(max_tx_bytes).to_bytes()],
+        };
+        match shell.process_proposal(request) {
+            Ok(_) => panic!("Test failed"),
+            Err(TestError::RejectProposal(response)) => {
+                assert_eq!(
+                    response[0].result.code,
+                    u32::from(ErrorCodes::TooLarge)
+                );
+            }
+        }
+
+        let request = ProcessProposal {
+            txs: vec![new_tx(0).to_bytes()],
+        };
+        match shell.process_proposal(request) {
+            Ok(_) => panic!("Test failed"),
+            Err(TestError::RejectProposal(response)) => {
+                assert!(
+                    response[0].result.code != u32::from(ErrorCodes::TooLarge)
+                );
+            }
+        }
+    }
 }