Fix channel closure and account tracking in callback package.

Before this change:

1. The account map wasn't populated if no callback was requested
   during execution.
2. Sequence number wasn't namespaced by channel_id, which would cause
   incorrect callbacks to be sent in the event of a channel closing.

contracts/note/src/contract.rs

@@ -4,13 +4,13 @@ use cosmwasm_std::{
     to_binary, Binary, Deps, DepsMut, Env, IbcMsg, IbcTimeout, MessageInfo, Response, StdResult,
 };
 use cw2::set_contract_version;
-use polytone::callback::CallbackRequestType;
-use polytone::{callback, ibc};
+use polytone::callbacks::CallbackRequestType;
+use polytone::{accounts, callbacks, ibc};
 
 use crate::error::ContractError;
 
 use crate::msg::{ExecuteMsg, InstantiateMsg, MigrateMsg, Pair, QueryMsg};
-use crate::state::{BLOCK_MAX_GAS, CHANNEL, CONNECTION_REMOTE_PORT};
+use crate::state::{increment_sequence_number, BLOCK_MAX_GAS, CHANNEL, CONNECTION_REMOTE_PORT};
 
 const CONTRACT_NAME: &str = "crates.io:polytone-note";
 const CONTRACT_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -74,17 +74,28 @@ pub fn execute(
         ),
     };
 
-    callback::request_callback(
+    let channel_id = CHANNEL
+        .may_load(deps.storage)?
+        .ok_or(ContractError::NoPair)?;
+
+    let sequence_number = increment_sequence_number(deps.storage, channel_id.clone())?;
+
+    callbacks::request_callback(
         deps.storage,
         deps.api,
+        channel_id.clone(),
+        sequence_number,
         info.sender.clone(),
         callback,
         request_type,
     )?;
 
-    let channel_id = CHANNEL
-        .may_load(deps.storage)?
-        .ok_or(ContractError::NoPair)?;
+    accounts::on_send_packet(
+        deps.storage,
+        channel_id.clone(),
+        sequence_number,
+        &info.sender,
+    )?;
 
     Ok(Response::default()
         .add_attribute("method", "execute")
@@ -109,10 +120,10 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> StdResult<Binary> {
                 remote_port,
             },
         )),
-        QueryMsg::RemoteAddress { local_address } => to_binary(
-            &callback::LOCAL_TO_REMOTE_ACCOUNT
-                .may_load(deps.storage, &deps.api.addr_validate(&local_address)?)?,
-        ),
+        QueryMsg::RemoteAddress { local_address } => to_binary(&accounts::query_account(
+            deps.storage,
+            deps.api.addr_validate(&local_address)?,
+        )?),
         QueryMsg::BlockMaxGas => to_binary(&BLOCK_MAX_GAS.load(deps.storage)?),
     }
 }

contracts/note/src/ibc.rs

@@ -5,7 +5,7 @@ use cosmwasm_std::{
     IbcChannelOpenResponse, IbcPacketAckMsg, IbcPacketReceiveMsg, IbcPacketTimeoutMsg,
     IbcReceiveResponse, Never, Reply, Response, SubMsg,
 };
-use polytone::{callback, handshake::note};
+use polytone::{accounts, callbacks, handshake::note};
 
 use crate::{
     error::ContractError,
@@ -94,15 +94,23 @@ pub fn ibc_packet_ack(
     _env: Env,
     ack: IbcPacketAckMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
-    let callback = callback::on_ack(deps.storage, &ack).map(|cosmos_msg| {
-        SubMsg::reply_on_error(cosmos_msg, ack.original_packet.sequence).with_gas_limit(
+    let (callback, executed_by) = callbacks::on_ack(deps.storage, &ack);
+    let callback = callback.map(|callback| {
+        SubMsg::reply_on_error(callback, ack.original_packet.sequence).with_gas_limit(
             BLOCK_MAX_GAS
                 .load(deps.storage)
                 .expect("set during instantiation")
                 - ERR_GAS_NEEDED,
         )
     });
 
+    accounts::on_ack(
+        deps.storage,
+        ack.original_packet.src.channel_id.clone(),
+        ack.original_packet.sequence,
+        executed_by,
+    );
+
     Ok(IbcBasicResponse::default()
         .add_attribute("method", "ibc_packet_ack")
         .add_attribute("sequence_number", ack.original_packet.sequence.to_string())
@@ -115,7 +123,7 @@ pub fn ibc_packet_timeout(
     _env: Env,
     msg: IbcPacketTimeoutMsg,
 ) -> Result<IbcBasicResponse, ContractError> {
-    let callback = callback::on_timeout(deps.storage, &msg).map(|cosmos_msg| {
+    let callback = callbacks::on_timeout(deps.storage, &msg).map(|cosmos_msg| {
         SubMsg::reply_on_error(cosmos_msg, msg.packet.sequence).with_gas_limit(
             BLOCK_MAX_GAS
                 .load(deps.storage)
@@ -124,6 +132,8 @@ pub fn ibc_packet_timeout(
         )
     });
 
+    accounts::on_timeout(deps.storage, msg.packet.src.channel_id, msg.packet.sequence);
+
     Ok(IbcBasicResponse::default()
         .add_attribute("method", "ibc_packet_timeout")
         .add_attribute("sequence_number", msg.packet.sequence.to_string())

contracts/note/src/msg.rs

@@ -1,7 +1,7 @@
 use cosmwasm_schema::{cw_serde, QueryResponses};
 use cosmwasm_std::{CosmosMsg, Empty, QueryRequest, Uint64};
 
-use polytone::callback::CallbackRequest;
+use polytone::callbacks::CallbackRequest;
 
 #[cw_serde]
 pub struct InstantiateMsg {

contracts/note/src/state.rs

@@ -1,4 +1,5 @@
-use cw_storage_plus::Item;
+use cosmwasm_std::{StdResult, Storage};
+use cw_storage_plus::{Item, Map};
 
 /// (Connection-ID, Remote port) of this contract's pair.
 pub const CONNECTION_REMOTE_PORT: Item<(String, String)> = Item::new("a");
@@ -8,4 +9,23 @@ pub const CONNECTION_REMOTE_PORT: Item<(String, String)> = Item::new("a");
 pub const CHANNEL: Item<String> = Item::new("b");
 
 /// Max gas usable in a single block.
-pub(crate) const BLOCK_MAX_GAS: Item<u64> = Item::new("bmg");
+pub const BLOCK_MAX_GAS: Item<u64> = Item::new("bmg");
+
+/// (channel_id) -> sequence number. `u64` is the type used in the
+/// Cosmos SDK for sequence numbers:
+///
+/// <https://github.com/cosmos/ibc-go/blob/a25f0d421c32b3a2b7e8168c9f030849797ff2e8/modules/core/02-client/keeper/keeper.go#L116-L125>
+const SEQUENCE_NUMBER: Map<String, u64> = Map::new("sn");
+
+/// Increments and returns the next sequence number.
+pub(crate) fn increment_sequence_number(
+    storage: &mut dyn Storage,
+    channel_id: String,
+) -> StdResult<u64> {
+    let seq = SEQUENCE_NUMBER
+        .may_load(storage, channel_id.clone())?
+        .unwrap_or_default()
+        + 1;
+    SEQUENCE_NUMBER.save(storage, channel_id, &seq)?;
+    Ok(seq)
+}

contracts/voice/src/ibc.rs

@@ -10,7 +10,7 @@ use cosmwasm_std::{
 use cw_utils::{parse_reply_execute_data, MsgExecuteContractResponse};
 use polytone::{
     ack::{ack_execute_fail, ack_fail},
-    callback::Callback,
+    callbacks::Callback,
     handshake::voice,
 };
 

packages/polytone/src/accounts.rs

@@ -0,0 +1,48 @@
+use cosmwasm_std::{Addr, StdResult, Storage};
+use cw_storage_plus::Map;
+
+/// (channel_id, sequence_number) -> sender
+///
+/// Maps packets to the address that sent them.
+const PENDING: Map<(String, u64), Addr> = Map::new("polytone-accounts-pending");
+
+/// (local_account) -> remote_account
+///
+/// Maps local addresses to their remote counterparts.
+const LOCAL_TO_REMOTE_ACCOUNT: Map<Addr, String> = Map::new("polytone-account-map");
+
+pub fn on_send_packet(
+    storage: &mut dyn Storage,
+    channel_id: String,
+    sequence_number: u64,
+    sender: &Addr,
+) -> StdResult<()> {
+    PENDING.save(storage, (channel_id, sequence_number), sender)
+}
+
+pub fn on_ack(
+    storage: &mut dyn Storage,
+    channel_id: String,
+    sequence_number: u64,
+    executor: Option<String>,
+) {
+    let local_account = PENDING
+        .load(storage, (channel_id.clone(), sequence_number))
+        .expect("pending was set when sending packet");
+
+    PENDING.remove(storage, (channel_id, sequence_number));
+
+    if let Some(executor) = executor {
+        LOCAL_TO_REMOTE_ACCOUNT
+            .save(storage, local_account, &executor)
+            .expect("strings were loaded from storage, so should serialize");
+    }
+}
+
+pub fn on_timeout(storage: &mut dyn Storage, channel_id: String, sequence_number: u64) {
+    PENDING.remove(storage, (channel_id, sequence_number))
+}
+
+pub fn query_account(storage: &dyn Storage, local_address: Addr) -> StdResult<Option<String>> {
+    LOCAL_TO_REMOTE_ACCOUNT.may_load(storage, local_address)
+}

packages/polytone/src/ack.rs

@@ -1,7 +1,7 @@
 use cosmwasm_std::{from_binary, to_binary, Binary, IbcAcknowledgement, SubMsgResponse, Uint64};
 
-pub use crate::callback::Callback;
-use crate::callback::{ErrorResponse, ExecutionResponse};
+pub use crate::callbacks::Callback;
+use crate::callbacks::{ErrorResponse, ExecutionResponse};
 
 /// wasmd 0.32+ will not return a hardcoded ICS-20 ACK if
 /// ibc_packet_receive errors [1] so we can safely use an ACK format