Merge pull request #411 from Concordium/fix-counter-notify

Fix counter notify
Concordium · Apr 15, 2024 · 03f677d · 03f677d
2 parents 49d9c78 + 607c80b
commit 03f677d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 14 deletions.
diff --git a/.github/workflows/linter.yml b/.github/workflows/linter.yml
@@ -26,34 +26,35 @@ jobs:
           - concordium-std-derive/Cargo.toml
           - concordium-std/Cargo.toml
           - concordium-cis2/Cargo.toml
-          - examples/voting/Cargo.toml
-          - examples/eSealing/Cargo.toml
+          - examples/account-signature-checks/Cargo.toml
           - examples/auction/Cargo.toml
+          - examples/bump-alloc-tests/Cargo.toml
           - examples/cis2-dynamic-nft/Cargo.toml
           - examples/cis2-multi/Cargo.toml
           - examples/cis2-multi-royalties/Cargo.toml
           - examples/cis2-nft/Cargo.toml
-          - examples/cis3-nft-sponsored-txs/Cargo.toml
           - examples/cis2-wccd/Cargo.toml
+          - examples/cis3-nft-sponsored-txs/Cargo.toml
+          - examples/counter-notify/Cargo.toml
           - examples/credential-registry/Cargo.toml
+          - examples/eSealing/Cargo.toml
           - examples/factory/Cargo.toml
           - examples/fib/Cargo.toml
           - examples/icecream/Cargo.toml
           - examples/memo/Cargo.toml
           - examples/nametoken/Cargo.toml
+          - examples/offchain-transfers/Cargo.toml
           - examples/piggy-bank/part1/Cargo.toml
           - examples/piggy-bank/part2/Cargo.toml
           - examples/proxy/Cargo.toml
           - examples/recorder/Cargo.toml
           - examples/signature-verifier/Cargo.toml
-          - examples/transfer-policy-check/Cargo.toml
-          - examples/two-step-transfer/Cargo.toml
           - examples/smart-contract-upgrade/contract-version1/Cargo.toml
           - examples/smart-contract-upgrade/contract-version2/Cargo.toml
-          - examples/offchain-transfers/Cargo.toml
-          - examples/account-signature-checks/Cargo.toml
           - examples/sponsored-tx-enabled-auction/Cargo.toml
-          - examples/bump-alloc-tests/Cargo.toml
+          - examples/transfer-policy-check/Cargo.toml
+          - examples/two-step-transfer/Cargo.toml
+          - examples/voting/Cargo.toml
 
     steps:
       - name: Checkout sources

diff --git a/examples/counter-notify/src/lib.rs b/examples/counter-notify/src/lib.rs
@@ -13,6 +13,12 @@ use concordium_std::*;
 
 type State = u64;
 
+#[derive(Serial, Deserial, PartialEq)]
+pub enum ReentryOccurance {
+    NoReentryAttack,
+    ReentryAttack,
+}
+
 #[init(contract = "counter-notify")]
 #[inline(always)]
 fn contract_init(_ctx: &InitContext, _state_builder: &mut StateBuilder) -> InitResult<State> {
@@ -29,9 +35,13 @@ fn just_increment(_ctx: &ReceiveContext, host: &mut Host<State>) -> ReceiveResul
     contract = "counter-notify",
     name = "increment-and-notify",
     mutable,
-    parameter = "(ContractAddress, OwnedEntrypointName)"
+    parameter = "(ContractAddress, OwnedEntrypointName)",
+    return_value = "ReentryOccurance"
 )]
-fn increment_and_notify(ctx: &ReceiveContext, host: &mut Host<State>) -> ReceiveResult<bool> {
+fn increment_and_notify(
+    ctx: &ReceiveContext,
+    host: &mut Host<State>,
+) -> ReceiveResult<ReentryOccurance> {
     let (contract, entrypoint): (ContractAddress, OwnedEntrypointName) =
         ctx.parameter_cursor().get()?;
 
@@ -48,8 +58,13 @@ fn increment_and_notify(ctx: &ReceiveContext, host: &mut Host<State>) -> Receive
     )
     .unwrap_abort();
 
-    // This will be false if the reentrancy has occurred.
-    Ok(preinvoke_count == *host.state())
+    let is_reentry = if preinvoke_count != *host.state() {
+        ReentryOccurance::ReentryAttack
+    } else {
+        ReentryOccurance::NoReentryAttack
+    };
+
+    Ok(is_reentry)
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////

diff --git a/examples/counter-notify/tests/tests.rs b/examples/counter-notify/tests/tests.rs
@@ -1,5 +1,6 @@
 use concordium_smart_contract_testing::*;
 use concordium_std_derive::account_address;
+use counter_notify::ReentryOccurance;
 
 const ACC_0: AccountAddress =
     account_address!("2wkBET2rRgE8pahuaczxKbmv7ciehqsne57F9gtzf1PVdr2VP3");
@@ -112,6 +113,6 @@ fn tests() {
     ] if rcv_name_1 == "counter-notify.just-increment" && rcv_name_2 == "reentrancy-attacker.call-just-increment" && rcv_name_3 == "counter-notify.increment-and-notify"));
 
     // Check that the contract observed the reentrancy attack.
-    let rv: bool = update.parse_return_value().unwrap();
-    assert!(rv, "Re-entrancy attack not observed.");
+    let rv: ReentryOccurance = update.parse_return_value().unwrap();
+    assert!(rv == ReentryOccurance::ReentryAttack, "Re-entrancy attack not observed.");
 }