From bf5068da7d93a7b65df669d360d548021f7218ce Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Mon, 17 Oct 2022 11:51:00 +0200
Subject: [PATCH] contract-sdk/specs/token/oas20: Rollback if subcall fails in
 send

---
 contract-sdk/specs/token/oas20/src/helpers.rs | 34 +++++++++++++++++--
 contract-sdk/specs/token/oas20/src/lib.rs     |  6 +++-
 contract-sdk/specs/token/oas20/src/types.rs   |  4 +++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/contract-sdk/specs/token/oas20/src/helpers.rs b/contract-sdk/specs/token/oas20/src/helpers.rs
index f8c3c528d2..7c30b39b34 100644
--- a/contract-sdk/specs/token/oas20/src/helpers.rs
+++ b/contract-sdk/specs/token/oas20/src/helpers.rs
@@ -2,7 +2,7 @@ use oasis_contract_sdk::{
     self as sdk,
     env::Env,
     types::{
-        message::{Message, NotifyReply},
+        message::{CallResult, Message, NotifyReply, Reply},
         InstanceId,
     },
 };
@@ -14,6 +14,9 @@ use crate::types::{
     TokenInstantiation,
 };
 
+/// Unique identifier for the send subcall.
+pub const CALL_ID_SEND: u64 = 1;
+
 /// Handles an OAS20 request call.
 pub fn handle_call<C: sdk::Context>(
     ctx: &mut C,
@@ -34,7 +37,16 @@ pub fn handle_call<C: sdk::Context>(
         }
         Request::Send { to, amount, data } => {
             let from = ctx.caller_address().to_owned();
-            send(ctx, balances, from, to, amount, data, 0, NotifyReply::Never)?;
+            send(
+                ctx,
+                balances,
+                from,
+                to,
+                amount,
+                data,
+                CALL_ID_SEND,
+                NotifyReply::OnError, // Rollback if subcall fails.
+            )?;
 
             ctx.emit_event(Event::Oas20Sent { from, to, amount });
 
@@ -120,6 +132,24 @@ pub fn handle_query<C: sdk::Context>(
     }
 }
 
+/// Handles a reply from OAS20 execution.
+pub fn handle_reply<C: sdk::Context>(
+    _ctx: &mut C,
+    _token_info: PublicCell<TokenInformation>,
+    _balances: PublicMap<Address, u128>,
+    _allowances: PublicMap<(Address, Address), u128>,
+    reply: Reply,
+) -> Result<Option<Response>, Error> {
+    match reply {
+        Reply::Call {
+            id: CALL_ID_SEND,
+            result: CallResult::Failed { module, code },
+            ..
+        } => Err(Error::ReceiverCallFailed(module, code)),
+        _ => Ok(None),
+    }
+}
+
 /// Instantiates the contract state.
 pub fn instantiate<C: sdk::Context>(
     ctx: &mut C,
diff --git a/contract-sdk/specs/token/oas20/src/lib.rs b/contract-sdk/specs/token/oas20/src/lib.rs
index 135a3d8e67..fa49273477 100644
--- a/contract-sdk/specs/token/oas20/src/lib.rs
+++ b/contract-sdk/specs/token/oas20/src/lib.rs
@@ -4,7 +4,7 @@ extern crate alloc;
 pub mod helpers;
 pub mod types;
 
-use oasis_contract_sdk::{self as sdk};
+use oasis_contract_sdk::{self as sdk, types::message::Reply};
 use oasis_contract_sdk_storage::{cell::PublicCell, map::PublicMap};
 use oasis_contract_sdk_types::address::Address;
 
@@ -50,6 +50,10 @@ impl sdk::Contract for Oas20Token {
     fn query<C: sdk::Context>(ctx: &mut C, request: Request) -> Result<Response, Error> {
         helpers::handle_query(ctx, TOKEN_INFO, BALANCES, ALLOWANCES, request)
     }
+
+    fn handle_reply<C: sdk::Context>(ctx: &mut C, reply: Reply) -> Result<Option<Response>, Error> {
+        helpers::handle_reply(ctx, TOKEN_INFO, BALANCES, ALLOWANCES, reply)
+    }
 }
 
 // Create the required WASM exports required for the contract to be runnable.
diff --git a/contract-sdk/specs/token/oas20/src/types.rs b/contract-sdk/specs/token/oas20/src/types.rs
index 14ea81a6e1..297ed151d5 100644
--- a/contract-sdk/specs/token/oas20/src/types.rs
+++ b/contract-sdk/specs/token/oas20/src/types.rs
@@ -90,6 +90,10 @@ pub enum Error {
     #[error("insufficient allowance")]
     #[sdk_error(code = 8)]
     InsufficientAllowance,
+
+    #[error("receiver call failed (module: {0} code: {1})")]
+    #[sdk_error(code = 9)]
+    ReceiverCallFailed(String, u32),
 }
 
 /// All possible events that can be returned by the OAS20 contract.