From 1d8e5677c7f981c9f44cc481871751f5a935e4ae Mon Sep 17 00:00:00 2001 From: ekez Date: Thu, 27 Apr 2023 18:15:23 -0700 Subject: [PATCH] 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 | 33 +++-- contracts/note/src/ibc.rs | 18 ++- contracts/note/src/msg.rs | 2 +- contracts/note/src/state.rs | 24 +++- contracts/voice/src/ibc.rs | 2 +- packages/polytone/src/accounts.rs | 48 +++++++ packages/polytone/src/ack.rs | 4 +- .../src/{callback.rs => callbacks.rs} | 131 ++++++++++-------- packages/polytone/src/lib.rs | 3 +- tests/polytone-tester/src/msg.rs | 4 +- tests/polytone-tester/src/state.rs | 2 +- 11 files changed, 185 insertions(+), 86 deletions(-) create mode 100644 packages/polytone/src/accounts.rs rename packages/polytone/src/{callback.rs => callbacks.rs} (76%) diff --git a/contracts/note/src/contract.rs b/contracts/note/src/contract.rs index d4a558b..0a75bbd 100644 --- a/contracts/note/src/contract.rs +++ b/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 { 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)?), } } diff --git a/contracts/note/src/ibc.rs b/contracts/note/src/ibc.rs index 86fa930..9d66e72 100644 --- a/contracts/note/src/ibc.rs +++ b/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,8 +94,9 @@ pub fn ibc_packet_ack( _env: Env, ack: IbcPacketAckMsg, ) -> Result { - 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") @@ -103,6 +104,13 @@ pub fn ibc_packet_ack( ) }); + 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 { - 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()) diff --git a/contracts/note/src/msg.rs b/contracts/note/src/msg.rs index 37567bd..6cf233d 100644 --- a/contracts/note/src/msg.rs +++ b/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 { diff --git a/contracts/note/src/state.rs b/contracts/note/src/state.rs index 711a20d..207831c 100644 --- a/contracts/note/src/state.rs +++ b/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 = Item::new("b"); /// Max gas usable in a single block. -pub(crate) const BLOCK_MAX_GAS: Item = Item::new("bmg"); +pub const BLOCK_MAX_GAS: Item = Item::new("bmg"); + +/// (channel_id) -> sequence number. `u64` is the type used in the +/// Cosmos SDK for sequence numbers: +/// +/// +const SEQUENCE_NUMBER: Map = Map::new("sn"); + +/// Increments and returns the next sequence number. +pub(crate) fn increment_sequence_number( + storage: &mut dyn Storage, + channel_id: String, +) -> StdResult { + let seq = SEQUENCE_NUMBER + .may_load(storage, channel_id.clone())? + .unwrap_or_default() + + 1; + SEQUENCE_NUMBER.save(storage, channel_id, &seq)?; + Ok(seq) +} diff --git a/contracts/voice/src/ibc.rs b/contracts/voice/src/ibc.rs index 2613c19..54bbe82 100644 --- a/contracts/voice/src/ibc.rs +++ b/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, }; diff --git a/packages/polytone/src/accounts.rs b/packages/polytone/src/accounts.rs new file mode 100644 index 0000000..9aba7eb --- /dev/null +++ b/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 = 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, +) { + 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> { + LOCAL_TO_REMOTE_ACCOUNT.may_load(storage, local_address) +} diff --git a/packages/polytone/src/ack.rs b/packages/polytone/src/ack.rs index 5c870cb..60382ba 100644 --- a/packages/polytone/src/ack.rs +++ b/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 diff --git a/packages/polytone/src/callback.rs b/packages/polytone/src/callbacks.rs similarity index 76% rename from packages/polytone/src/callback.rs rename to packages/polytone/src/callbacks.rs index 0db5cf9..ae91782 100644 --- a/packages/polytone/src/callback.rs +++ b/packages/polytone/src/callbacks.rs @@ -3,7 +3,7 @@ use cosmwasm_std::{ to_binary, Addr, Api, Binary, CosmosMsg, IbcPacketAckMsg, IbcPacketTimeoutMsg, StdResult, Storage, SubMsgResponse, Uint64, WasmMsg, }; -use cw_storage_plus::{Item, Map}; +use cw_storage_plus::Map; use crate::ack::unmarshal_ack; @@ -82,28 +82,32 @@ pub struct CallbackRequest { pub msg: Binary, } -//// Must be called every time a packet is sent to keep sequence -//// number tracking accurate. -/// -/// Requests that a callback be returned for the next IBC message that -/// is sent. +/// Disembiguates between a callback for remote message execution and +/// queries. +#[cw_serde] +pub enum CallbackRequestType { + Execute, + Query, +} + +/// Requests that a callback be returned for the IBC message +/// identified by `(channel_id, sequence_number)`. pub fn request_callback( storage: &mut dyn Storage, api: &dyn Api, + channel_id: String, + sequence_number: u64, initiator: Addr, request: Option, request_type: CallbackRequestType, ) -> StdResult<()> { - let seq = SEQ.may_load(storage)?.unwrap_or_default() + 1; - SEQ.save(storage, &seq)?; - if let Some(request) = request { let receiver = api.addr_validate(&request.receiver)?; let initiator_msg = request.msg; CALLBACKS.save( storage, - seq, + (channel_id, sequence_number), &PendingCallback { initiator, initiator_msg, @@ -116,28 +120,11 @@ pub fn request_callback( Ok(()) } -fn callback_msg(request: PendingCallback, result: Callback) -> CosmosMsg { - /// Gives the executed message a "callback" tag: - /// `{ "callback": CallbackMsg }`. - #[cw_serde] - enum C { - Callback(CallbackMessage), - } - WasmMsg::Execute { - contract_addr: request.receiver.into_string(), - msg: to_binary(&C::Callback(CallbackMessage { - initiator: request.initiator, - initiator_msg: request.initiator_msg, - result, - })) - .expect("fields are known to be serializable"), - funds: vec![], - } - .into() -} - /// Call on every packet ACK. Returns a callback message to execute, -/// if any. +/// if any, and the address that executed the request on the remote +/// chain (the message initiator's remote account), if any. +/// +/// (storage, ack) -> (callback, executed_by) pub fn on_ack( storage: &mut dyn Storage, IbcPacketAckMsg { @@ -145,18 +132,23 @@ pub fn on_ack( original_packet, .. }: &IbcPacketAckMsg, -) -> Option { - let Some(request) = CALLBACKS.may_load(storage, original_packet.sequence).unwrap() else { - return None - }; - CALLBACKS.remove(storage, original_packet.sequence); +) -> (Option, Option) { let result = unmarshal_ack(acknowledgement); - if let Callback::Execute(Ok(ExecutionResponse { executed_by, .. })) = &result { - LOCAL_TO_REMOTE_ACCOUNT - .save(storage, &request.initiator, executed_by) - .expect("strings can be serialized"); - } - Some(callback_msg(request, result)) + + let executed_by = match result { + Callback::Execute(Ok(ExecutionResponse { + ref executed_by, .. + })) => Some(executed_by.clone()), + _ => None, + }; + let callback_message = dequeue_callback( + storage, + original_packet.src.channel_id.clone(), + original_packet.sequence, + ) + .map(|request| callback_message(request, result)); + + (callback_message, executed_by) } /// Call on every packet timeout. Returns a callback message to execute, @@ -165,10 +157,7 @@ pub fn on_timeout( storage: &mut dyn Storage, IbcPacketTimeoutMsg { packet, .. }: &IbcPacketTimeoutMsg, ) -> Option { - let Some(request) = CALLBACKS.may_load(storage, packet.sequence).unwrap() else { - return None - }; - CALLBACKS.remove(storage, packet.sequence); + let request = dequeue_callback(storage, packet.src.channel_id.clone(), packet.sequence)?; let timeout = "timeout".to_string(); let result = match request.request_type { CallbackRequestType::Execute => Callback::Execute(Err(timeout)), @@ -177,7 +166,39 @@ pub fn on_timeout( error: timeout, })), }; - Some(callback_msg(request, result)) + Some(callback_message(request, result)) +} + +fn callback_message(request: PendingCallback, result: Callback) -> CosmosMsg { + /// Gives the executed message a "callback" tag: + /// `{ "callback": CallbackMsg }`. + #[cw_serde] + enum C { + Callback(CallbackMessage), + } + WasmMsg::Execute { + contract_addr: request.receiver.into_string(), + msg: to_binary(&C::Callback(CallbackMessage { + initiator: request.initiator, + initiator_msg: request.initiator_msg, + result, + })) + .expect("fields are known to be serializable"), + funds: vec![], + } + .into() +} + +fn dequeue_callback( + storage: &mut dyn Storage, + channel_id: String, + sequence_number: u64, +) -> Option { + let request = CALLBACKS + .may_load(storage, (channel_id.clone(), sequence_number)) + .unwrap()?; + CALLBACKS.remove(storage, (channel_id, sequence_number)); + Some(request) } #[cw_serde] @@ -190,17 +211,5 @@ struct PendingCallback { request_type: CallbackRequestType, } -/// Disembiguates between a callback for remote message execution and -/// queries. -#[cw_serde] -pub enum CallbackRequestType { - Execute, - Query, -} - -/// (local_account) -> remote_account -pub const LOCAL_TO_REMOTE_ACCOUNT: Map<&Addr, String> = Map::new("polytone-l2r"); -/// (sequence_number) -> callback -const CALLBACKS: Map = Map::new("polytone-callbacks"); -/// The number of packets sent so far. -const SEQ: Item = Item::new("polytone-ibc-seq"); +/// (channel_id, sequence_number) -> callback +const CALLBACKS: Map<(String, u64), PendingCallback> = Map::new("polytone-callbacks"); diff --git a/packages/polytone/src/lib.rs b/packages/polytone/src/lib.rs index dbbe2c5..1184f95 100644 --- a/packages/polytone/src/lib.rs +++ b/packages/polytone/src/lib.rs @@ -1,5 +1,6 @@ +pub mod accounts; pub mod ack; -pub mod callback; +pub mod callbacks; pub mod ibc; pub mod handshake; diff --git a/tests/polytone-tester/src/msg.rs b/tests/polytone-tester/src/msg.rs index b85850c..bf7bff1 100644 --- a/tests/polytone-tester/src/msg.rs +++ b/tests/polytone-tester/src/msg.rs @@ -9,7 +9,7 @@ pub enum ExecuteMsg { /// Calls `set_data(data)` if `data` is not None. Hello { data: Option }, /// Stores the callback in state and makes it queryable - Callback(polytone::callback::CallbackMessage), + Callback(polytone::callbacks::CallbackMessage), /// Runs out of gas. RunOutOfGas {}, } @@ -28,7 +28,7 @@ pub enum QueryMsg { #[cw_serde] pub struct CallbackHistoryResponse { - pub history: Vec, + pub history: Vec, } #[cw_serde] diff --git a/tests/polytone-tester/src/state.rs b/tests/polytone-tester/src/state.rs index 07f8a3b..a4c7314 100644 --- a/tests/polytone-tester/src/state.rs +++ b/tests/polytone-tester/src/state.rs @@ -1,5 +1,5 @@ use cw_storage_plus::Item; -use polytone::callback::CallbackMessage; +use polytone::callbacks::CallbackMessage; pub(crate) const CALLBACK_HISTORY: Item> = Item::new("a"); pub(crate) const HELLO_CALL_HISTORY: Item> = Item::new("b");