From 859d5b7bc5a700fa77e46bac600d018185ff2b46 Mon Sep 17 00:00:00 2001
From: ptrus <peter@u-s.si>
Date: Tue, 6 Jun 2023 15:18:16 +0200
Subject: [PATCH] runtime-sdk: support aborting in `after_handle_call`

---
 runtime-sdk/src/dispatcher.rs            |  8 +++++++-
 runtime-sdk/src/module.rs                | 16 +++++++++++++---
 runtime-sdk/src/modules/accounts/mod.rs  |  7 ++++++-
 runtime-sdk/src/modules/accounts/test.rs |  5 +++--
 runtime-sdk/src/modules/core/mod.rs      |  7 ++++++-
 runtime-sdk/src/modules/core/test.rs     |  5 +++--
 6 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/runtime-sdk/src/dispatcher.rs b/runtime-sdk/src/dispatcher.rs
index ac285068be..80099c1093 100644
--- a/runtime-sdk/src/dispatcher.rs
+++ b/runtime-sdk/src/dispatcher.rs
@@ -196,7 +196,13 @@ impl<R: Runtime> Dispatcher<R> {
         let (result, metadata) = Self::_dispatch_tx_call(ctx, call, opts);
 
         // Unconditionally call after handle call hook.
-        R::Modules::after_handle_call(ctx, &result);
+        let result = match R::Modules::after_handle_call(ctx, result) {
+            Ok(result) => result,
+            Err(e) => {
+                // If the call failed, return the error.
+                return (e.into_call_result(), metadata);
+            }
+        };
 
         // Make sure that a read-only call did not result in any modifications.
         if read_only && ctx.runtime_state().has_pending_updates() {
diff --git a/runtime-sdk/src/module.rs b/runtime-sdk/src/module.rs
index ea67ece989..a6dc25f5a4 100644
--- a/runtime-sdk/src/module.rs
+++ b/runtime-sdk/src/module.rs
@@ -382,8 +382,12 @@ pub trait TransactionHandler {
     /// Perform any action after call, within the transaction context.
     ///
     /// If an error is returned the transaction call fails and updates are rolled back.
-    fn after_handle_call<C: TxContext>(_ctx: &mut C, _result: &CallResult) {
+    fn after_handle_call<C: TxContext>(
+        _ctx: &mut C,
+        result: CallResult,
+    ) -> Result<CallResult, modules::core::Error> {
         // Default implementation doesn't do anything.
+        Ok(result)
     }
 
     /// Perform any action after dispatching the transaction, in batch context.
@@ -437,8 +441,14 @@ impl TransactionHandler for Tuple {
         Ok(())
     }
 
-    fn after_handle_call<C: TxContext>(ctx: &mut C, result: &CallResult) {
-        for_tuples!( #( Tuple::after_handle_call(ctx, result); )* );
+    fn after_handle_call<C: TxContext>(
+        ctx: &mut C,
+        mut result: CallResult,
+    ) -> Result<CallResult, modules::core::Error> {
+        for_tuples!( #(
+            result = Tuple::after_handle_call(ctx, result)?;
+        )* );
+        Ok(result)
     }
 
     fn after_dispatch_tx<C: Context>(ctx: &mut C, tx_auth_info: &AuthInfo, result: &CallResult) {
diff --git a/runtime-sdk/src/modules/accounts/mod.rs b/runtime-sdk/src/modules/accounts/mod.rs
index 97645bcb14..aa9ca52280 100644
--- a/runtime-sdk/src/modules/accounts/mod.rs
+++ b/runtime-sdk/src/modules/accounts/mod.rs
@@ -952,7 +952,10 @@ impl module::TransactionHandler for Module {
         Ok(())
     }
 
-    fn after_handle_call<C: TxContext>(ctx: &mut C, _result: &module::CallResult) {
+    fn after_handle_call<C: TxContext>(
+        ctx: &mut C,
+        result: module::CallResult,
+    ) -> Result<module::CallResult, modules::core::Error> {
         let acc = ctx
             .value::<FeeAccumulator>(CONTEXT_KEY_TX_FEE_ACCUMULATOR)
             .take()
@@ -974,6 +977,8 @@ impl module::TransactionHandler for Module {
         for (denom, amount) in acc.total_fees.into_iter() {
             fee_acc.add(&token::BaseUnits(amount, denom));
         }
+
+        Ok(result)
     }
 
     fn after_dispatch_tx<C: Context>(
diff --git a/runtime-sdk/src/modules/accounts/test.rs b/runtime-sdk/src/modules/accounts/test.rs
index e105c42697..5bf576d14f 100644
--- a/runtime-sdk/src/modules/accounts/test.rs
+++ b/runtime-sdk/src/modules/accounts/test.rs
@@ -575,8 +575,9 @@ fn test_fee_disbursement() {
         // Run after call tx handler.
         Accounts::after_handle_call(
             &mut tx_ctx,
-            &module::CallResult::Ok(cbor::Value::Simple(cbor::SimpleValue::NullValue)),
-        );
+            module::CallResult::Ok(cbor::Value::Simple(cbor::SimpleValue::NullValue)),
+        )
+        .expect("after_handle_call should succeed");
         tx_ctx.commit()
     });
 
diff --git a/runtime-sdk/src/modules/core/mod.rs b/runtime-sdk/src/modules/core/mod.rs
index e2decba4d6..a6519e09c0 100644
--- a/runtime-sdk/src/modules/core/mod.rs
+++ b/runtime-sdk/src/modules/core/mod.rs
@@ -999,12 +999,17 @@ impl<Cfg: Config> module::TransactionHandler for Module<Cfg> {
         Ok(())
     }
 
-    fn after_handle_call<C: TxContext>(ctx: &mut C, _result: &module::CallResult) {
+    fn after_handle_call<C: TxContext>(
+        ctx: &mut C,
+        result: module::CallResult,
+    ) -> Result<module::CallResult, Error> {
         // Emit gas used event.
         if Cfg::EMIT_GAS_USED_EVENTS {
             let used_gas = Self::used_tx_gas(ctx);
             ctx.emit_unconditional_event(Event::GasUsed { amount: used_gas });
         }
+
+        Ok(result)
     }
 }
 
diff --git a/runtime-sdk/src/modules/core/test.rs b/runtime-sdk/src/modules/core/test.rs
index 20885f5ea6..de9f66e1de 100644
--- a/runtime-sdk/src/modules/core/test.rs
+++ b/runtime-sdk/src/modules/core/test.rs
@@ -1088,8 +1088,9 @@ fn test_gas_used_events() {
         assert_eq!(Core::used_tx_gas(&mut tx_ctx), 10);
         Core::after_handle_call(
             &mut tx_ctx,
-            &module::CallResult::Ok(cbor::Value::Simple(cbor::SimpleValue::NullValue)),
-        );
+            module::CallResult::Ok(cbor::Value::Simple(cbor::SimpleValue::NullValue)),
+        )
+        .expect("after_handle_call should succeed");
 
         let (etags, _) = tx_ctx.commit();
         let tags = etags.clone().into_tags();