From 2b44ac75668bb2aa70fa74c01c891310b03d5634 Mon Sep 17 00:00:00 2001
From: Christoph Otter <chris@confio.gmbh>
Date: Thu, 18 Apr 2024 17:05:00 +0200
Subject: [PATCH] Make source chain callbacks extendable

---
 contracts/ibc-callbacks/src/contract.rs | 25 +++++++++++++++++++------
 packages/std/src/ibc/callbacks.rs       | 12 +++++++++---
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/contracts/ibc-callbacks/src/contract.rs b/contracts/ibc-callbacks/src/contract.rs
index ebaec8e827..950405b4d8 100644
--- a/contracts/ibc-callbacks/src/contract.rs
+++ b/contracts/ibc-callbacks/src/contract.rs
@@ -1,8 +1,8 @@
 use cosmwasm_std::{
     entry_point, to_json_binary, to_json_string, Binary, Deps, DepsMut, Empty, Env,
     IbcBasicResponse, IbcCallbackRequest, IbcDestinationChainCallbackMsg, IbcDstCallback, IbcMsg,
-    IbcSourceChainCallbackMsg, IbcSrcCallback, IbcTimeout, MessageInfo, Response, StdError,
-    StdResult,
+    IbcPacketAckMsg, IbcPacketTimeoutMsg, IbcSourceChainCallbackMsg, IbcSrcCallback, IbcTimeout,
+    MessageInfo, Response, StdError, StdResult,
 };
 
 use crate::msg::{CallbackType, ExecuteMsg, QueryMsg};
@@ -82,13 +82,26 @@ pub fn ibc_source_chain_callback(
     let mut counts = load_stats(deps.storage)?;
 
     match msg {
-        IbcSourceChainCallbackMsg::Acknowledgement(ack) => {
+        IbcSourceChainCallbackMsg::Acknowledgement {
+            acknowledgement,
+            original_packet,
+            relayer,
+            ..
+        } => {
             // save the ack
-            counts.ibc_ack_callbacks.push(ack);
+            counts.ibc_ack_callbacks.push(IbcPacketAckMsg::new(
+                acknowledgement,
+                original_packet,
+                relayer,
+            ));
         }
-        IbcSourceChainCallbackMsg::Timeout(timeout) => {
+        IbcSourceChainCallbackMsg::Timeout {
+            packet, relayer, ..
+        } => {
             // save the timeout
-            counts.ibc_timeout_callbacks.push(timeout);
+            counts
+                .ibc_timeout_callbacks
+                .push(IbcPacketTimeoutMsg::new(packet, relayer));
         }
     }
 
diff --git a/packages/std/src/ibc/callbacks.rs b/packages/std/src/ibc/callbacks.rs
index fa253bc5a9..f83c1c1216 100644
--- a/packages/std/src/ibc/callbacks.rs
+++ b/packages/std/src/ibc/callbacks.rs
@@ -5,7 +5,7 @@ use cosmwasm_core::Binary;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
-use crate::{Addr, IbcPacket, IbcPacketAckMsg, IbcPacketTimeoutMsg, Uint64};
+use crate::{Addr, IbcAcknowledgement, IbcPacket, Uint64};
 
 /// This is just a type representing the data that has to be sent with the IBC message to receive
 /// callbacks. It should be serialized and sent with the IBC message.
@@ -108,8 +108,14 @@ pub struct IbcDstCallback {
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
 #[serde(rename_all = "snake_case")]
 pub enum IbcSourceChainCallbackMsg {
-    Acknowledgement(IbcPacketAckMsg),
-    Timeout(IbcPacketTimeoutMsg),
+    #[non_exhaustive]
+    Acknowledgement {
+        acknowledgement: IbcAcknowledgement,
+        original_packet: IbcPacket,
+        relayer: Addr,
+    },
+    #[non_exhaustive]
+    Timeout { packet: IbcPacket, relayer: Addr },
 }
 
 /// The message type of the IBC destination chain callback.