From 914b6f5bd37648304c1bb028ceb5f5b4dbf5341c Mon Sep 17 00:00:00 2001
From: Marco Granelli <marco@heliax.dev>
Date: Fri, 13 Oct 2023 20:00:34 +0200
Subject: [PATCH] Reworks replay protection check

---
 apps/src/lib/node/ledger/shell/mod.rs         | 39 ++++++++-----------
 .../lib/node/ledger/shell/prepare_proposal.rs |  9 +++--
 .../lib/node/ledger/shell/process_proposal.rs | 17 +++-----
 3 files changed, 29 insertions(+), 36 deletions(-)

diff --git a/apps/src/lib/node/ledger/shell/mod.rs b/apps/src/lib/node/ledger/shell/mod.rs
index 81aa225655..d81446baef 100644
--- a/apps/src/lib/node/ledger/shell/mod.rs
+++ b/apps/src/lib/node/ledger/shell/mod.rs
@@ -929,47 +929,42 @@ where
     pub fn replay_protection_checks(
         &self,
         wrapper: &Tx,
-        tx_bytes: &[u8],
         temp_wl_storage: &mut TempWlStorage<D, H>,
     ) -> Result<()> {
-        let inner_tx_hash =
-            wrapper.clone().update_header(TxType::Raw).header_hash();
+        let wrapper_hash = wrapper.header_hash();
         if temp_wl_storage
-            .has_replay_protection_entry(&inner_tx_hash)
-            .expect("Error while checking inner tx hash key in storage")
+            .has_replay_protection_entry(&wrapper_hash)
+            .expect("Error while checking wrapper tx hash key in storage")
         {
             return Err(Error::ReplayAttempt(format!(
-                "Inner transaction hash {} already in storage",
-                &inner_tx_hash,
+                "Wrapper transaction hash {} already in storage",
+                wrapper_hash
             )));
         }
 
-        // Write inner hash to tx WAL
+        // Write wrapper hash to tx WAL
         temp_wl_storage
             .write_log
-            .write_tx_hash(inner_tx_hash)
-            .expect("Couldn't write inner transaction hash to write log");
+            .write_tx_hash(wrapper_hash)
+            .map_err(|e| Error::ReplayAttempt(e.to_string()))?;
 
-        let tx =
-            Tx::try_from(tx_bytes).expect("Deserialization shouldn't fail");
-        let wrapper_hash = tx.header_hash();
+        let inner_tx_hash =
+            wrapper.clone().update_header(TxType::Raw).header_hash();
         if temp_wl_storage
-            .has_replay_protection_entry(&wrapper_hash)
-            .expect("Error while checking wrapper tx hash key in storage")
+            .has_replay_protection_entry(&inner_tx_hash)
+            .expect("Error while checking inner tx hash key in storage")
         {
             return Err(Error::ReplayAttempt(format!(
-                "Wrapper transaction hash {} already in storage",
-                wrapper_hash
+                "Inner transaction hash {} already in storage",
+                &inner_tx_hash,
             )));
         }
 
-        // Write wrapper hash to tx WAL
+        // Write inner hash to tx WAL
         temp_wl_storage
             .write_log
-            .write_tx_hash(wrapper_hash)
-            .expect("Couldn't write wrapper tx hash to write log");
-
-        Ok(())
+            .write_tx_hash(inner_tx_hash)
+            .map_err(|e| Error::ReplayAttempt(e.to_string()))
     }
 
     /// If a handle to an Ethereum oracle was provided to the [`Shell`], attempt
diff --git a/apps/src/lib/node/ledger/shell/prepare_proposal.rs b/apps/src/lib/node/ledger/shell/prepare_proposal.rs
index 103c90b982..20b9105875 100644
--- a/apps/src/lib/node/ledger/shell/prepare_proposal.rs
+++ b/apps/src/lib/node/ledger/shell/prepare_proposal.rs
@@ -245,8 +245,11 @@ where
             let mut tx_gas_meter = TxGasMeter::new(wrapper.gas_limit);
             tx_gas_meter.add_tx_size_gas(tx_bytes).map_err(|_| ())?;
 
-            // Check replay protection
-            self.replay_protection_checks(&tx, tx_bytes, temp_wl_storage)
+            // Check replay protection, safe to do here. Even if the tx is a
+            // replay attempt, we can leave its hashes in the write log since,
+            // having we already checked the signature, no other tx with the
+            // same hash can ba deemed valid
+            self.replay_protection_checks(&tx, temp_wl_storage)
                 .map_err(|_| ())?;
 
             // Check fees
@@ -1314,7 +1317,7 @@ mod test_prepare_proposal {
         let (shell, _recv, _, _) = test_utils::setup();
 
         let keypair = crate::wallet::defaults::daewon_keypair();
-        let keypair_2 = crate::wallet::defaults::daewon_keypair();
+        let keypair_2 = crate::wallet::defaults::albert_keypair();
         let mut wrapper =
             Tx::from_type(TxType::Wrapper(Box::new(WrapperTx::new(
                 Fee {
diff --git a/apps/src/lib/node/ledger/shell/process_proposal.rs b/apps/src/lib/node/ledger/shell/process_proposal.rs
index 50950b60f0..6f1c7c83b9 100644
--- a/apps/src/lib/node/ledger/shell/process_proposal.rs
+++ b/apps/src/lib/node/ledger/shell/process_proposal.rs
@@ -878,11 +878,9 @@ where
                     }
                 } else {
                     // Replay protection checks
-                    if let Err(e) = self.replay_protection_checks(
-                        &tx,
-                        tx_bytes,
-                        temp_wl_storage,
-                    ) {
+                    if let Err(e) =
+                        self.replay_protection_checks(&tx, temp_wl_storage)
+                    {
                         return TxResult {
                             code: ErrorCodes::ReplayTx.into(),
                             info: e.to_string(),
@@ -2227,12 +2225,9 @@ mod test_process_proposal {
                 assert_eq!(
                     response[1].result.info,
                     format!(
-                        "Transaction replay attempt: Inner transaction hash \
+                        "Transaction replay attempt: Wrapper transaction hash \
                          {} already in storage",
-                        wrapper
-                            .clone()
-                            .update_header(TxType::Raw)
-                            .header_hash(),
+                        wrapper.header_hash(),
                     )
                 );
             }
@@ -2311,7 +2306,7 @@ mod test_process_proposal {
         let (shell, _recv, _, _) = test_utils::setup();
 
         let keypair = crate::wallet::defaults::daewon_keypair();
-        let keypair_2 = crate::wallet::defaults::daewon_keypair();
+        let keypair_2 = crate::wallet::defaults::albert_keypair();
 
         let mut wrapper =
             Tx::from_type(TxType::Wrapper(Box::new(WrapperTx::new(