Factor out extracting tokens from IBC receive packet message into sep…

…arate function.
anoma · Feb 10, 2025 · ef4471d · ef4471d
1 parent af8d276
commit ef4471d
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 66 deletions.
diff --git a/crates/ibc/src/lib.rs b/crates/ibc/src/lib.rs
@@ -692,26 +692,22 @@ where
                         })?,
                 );
                 // Record the tokens credited/debited in this NFT transfer
-                let ibc_traces: Vec<_> = msg
+                let tokens = msg
                     .message
                     .packet_data
                     .token_ids
                     .0
                     .iter()
                     .map(|token_id| {
-                        ibc_trace_for_nft(
+                        convert_to_address(ibc_trace_for_nft(
                             &msg.message.packet_data.class_id,
                             token_id,
-                        )
-                    })
-                    .collect();
-                let mut tokens = BTreeSet::new();
-                for ibc_trace in ibc_traces {
-                    let token = convert_to_address(ibc_trace)
+                        ))
                         .into_storage_result()
-                        .map_err(Error::Storage)?;
-                    tokens.insert(token);
-                }
+                        .map_err(Error::Storage)
+                    })
+                    .collect::<Result<_, _>>()?;
+
                 send_nft_transfer_execute(
                     &mut self.ctx,
                     &mut nft_transfer_ctx,
@@ -746,69 +742,23 @@ where
                             .is_receiving_success(msg)?
                             .is_some_and(|ack_succ| ack_succ) =>
                     {
-                        let tokens = if msg.packet.port_id_on_b.as_str()
-                            == PORT_ID_STR
-                        {
-                            // Record the token credited/debited in this
-                            // transfer
-                            let packet_data =
-                                serde_json::from_slice::<PacketData>(
-                                    &msg.packet.data,
-                                )
-                                .map_err(StorageError::new)
-                                .map_err(Error::Storage)?;
-                            let ibc_denom = packet_data.token.denom.to_string();
+                        let ibc_traces = extract_traces_from_recv_msg(msg)
+                            .map_err(StorageError::new)
+                            .map_err(Error::Storage)?;
+                        let mut tokens = BTreeSet::new();
+                        for ibc_trace in ibc_traces {
                             // Get the received token
                             let token = received_ibc_token(
-                                ibc_denom,
+                                ibc_trace,
                                 &msg.packet.port_id_on_a,
                                 &msg.packet.chan_id_on_a,
                                 &msg.packet.port_id_on_b,
                                 &msg.packet.chan_id_on_b,
                             )
                             .into_storage_result()
                             .map_err(Error::Storage)?;
-                            [token].into()
-                        } else if msg.packet.port_id_on_b.as_str()
-                            == NFT_PORT_ID_STR
-                        {
-                            // Record the tokenS credited/debited in this NFT
-                            // transfer
-                            let packet_data =
-                                serde_json::from_slice::<NftPacketData>(
-                                    &msg.packet.data,
-                                )
-                                .map_err(StorageError::new)
-                                .map_err(Error::Storage)?;
-                            let ibc_traces: Vec<_> = packet_data
-                                .token_ids
-                                .0
-                                .iter()
-                                .map(|token_id| {
-                                    ibc_trace_for_nft(
-                                        &packet_data.class_id,
-                                        token_id,
-                                    )
-                                })
-                                .collect();
-                            let mut tokens = BTreeSet::new();
-                            for ibc_trace in ibc_traces {
-                                // Get the received token
-                                let token = received_ibc_token(
-                                    ibc_trace,
-                                    &msg.packet.port_id_on_a,
-                                    &msg.packet.chan_id_on_a,
-                                    &msg.packet.port_id_on_b,
-                                    &msg.packet.chan_id_on_b,
-                                )
-                                .into_storage_result()
-                                .map_err(Error::Storage)?;
-                                tokens.insert(token);
-                            }
-                            tokens
-                        } else {
-                            BTreeSet::new()
-                        };
+                            tokens.insert(token);
+                        }
                         (extract_masp_tx_from_packet(&msg.packet), tokens)
                     }
                     #[cfg(is_apple_silicon)]

diff --git a/crates/ibc/src/msg.rs b/crates/ibc/src/msg.rs
@@ -11,7 +11,9 @@ use ibc::apps::nft_transfer::types::PORT_ID_STR as NFT_PORT_ID_STR;
 use ibc::apps::transfer::types::msgs::transfer::MsgTransfer as IbcMsgTransfer;
 use ibc::apps::transfer::types::packet::PacketData;
 use ibc::apps::transfer::types::PORT_ID_STR as FT_PORT_ID_STR;
-use ibc::core::channel::types::msgs::PacketMsg;
+use ibc::core::channel::types::msgs::{
+    MsgRecvPacket as IbcMsgRecvPacket, PacketMsg,
+};
 use ibc::core::channel::types::packet::Packet;
 use ibc::core::handler::types::msgs::MsgEnvelope;
 use ibc::core::host::types::identifiers::PortId;
@@ -21,6 +23,8 @@ use namada_core::borsh::BorshSerializeExt;
 use namada_core::string_encoding::StringEncoded;
 use serde::{Deserialize, Serialize};
 
+use crate::trace;
+
 trait Sealed {}
 
 /// Marker trait that denotes whether an IBC memo is valid
@@ -327,6 +331,38 @@ fn extract_memo_from_packet(
     }
 }
 
+/// Extract the traces used in a receive packet message
+pub fn extract_traces_from_recv_msg(
+    msg: &IbcMsgRecvPacket,
+) -> Result<Vec<String>, serde_json::Error> {
+    match msg.packet.port_id_on_b.as_str() {
+        FT_PORT_ID_STR => {
+            // Record the token credited/debited in this
+            // transfer
+            let packet_data =
+                serde_json::from_slice::<PacketData>(&msg.packet.data)?;
+            let ibc_denom = packet_data.token.denom.to_string();
+            Ok([ibc_denom].into())
+        }
+        NFT_PORT_ID_STR => {
+            // Record the tokens credited/debited in this NFT
+            // transfer
+            let packet_data =
+                serde_json::from_slice::<NftPacketData>(&msg.packet.data)?;
+            let ibc_traces: Vec<_> = packet_data
+                .token_ids
+                .0
+                .iter()
+                .map(|token_id| {
+                    trace::ibc_trace_for_nft(&packet_data.class_id, token_id)
+                })
+                .collect();
+            Ok(ibc_traces)
+        }
+        _ => Ok(vec![]),
+    }
+}
+
 /// Get IBC memo string from MASP transaction for receiving
 pub fn convert_masp_tx_to_ibc_memo(transaction: &MaspTransaction) -> String {
     IbcShieldingData(transaction.clone()).into()