imp(Router)!: simplify Module lookup (#806)

* imp(router): simplify Module lookup

* fix: clippy catches

* imp: discard ModuleId and use PortId::transfer everywhere instead

* Revert "imp: discard ModuleId and use PortId::transfer everywhere instead"

This reverts commit 9227ef7.

* nit: RouterError variants

* fix: cargo test

.changelog/unreleased/breaking-changes/802-simplify-module-lookup.md

@@ -0,0 +1,2 @@
+- Simplify Module lookup in the `Router` trait
+ ([#802](https://github.com/cosmos/ibc-rs/issues/802))

crates/ibc/src/core/context.rs

@@ -33,6 +33,7 @@ use crate::Height;
 use super::ics02_client::client_state::ClientState;
 use super::ics02_client::consensus_state::ConsensusState;
 use super::ics02_client::ClientExecutionContext;
+use super::ics24_host::identifier::PortId;
 
 /// Top-level error
 #[derive(Debug, Display, From)]
@@ -69,6 +70,10 @@ pub enum RouterError {
     UnknownMessageTypeUrl { url: String },
     /// the message is malformed and cannot be decoded error: `{0}`
     MalformedMessageBytes(ibc_proto::protobuf::Error),
+    /// port `{port_id}` is unknown
+    UnknownPort { port_id: PortId },
+    /// module not found
+    ModuleNotFound,
 }
 
 impl From<ContextError> for RouterError {
@@ -82,8 +87,8 @@ impl std::error::Error for RouterError {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
         match &self {
             Self::ContextError(e) => Some(e),
-            Self::UnknownMessageTypeUrl { .. } => None,
             Self::MalformedMessageBytes(e) => Some(e),
+            _ => None,
         }
     }
 }

crates/ibc/src/core/handler.rs

@@ -5,7 +5,6 @@ use super::ics03_connection::handler::{
     conn_open_ack, conn_open_confirm, conn_open_init, conn_open_try,
 };
 use super::ics03_connection::msgs::ConnectionMsg;
-use super::ics04_channel::error::ChannelError;
 use super::ics04_channel::handler::acknowledgement::{
     acknowledgement_packet_execute, acknowledgement_packet_validate,
 };
@@ -27,10 +26,12 @@ use super::ics04_channel::handler::recv_packet::{recv_packet_execute, recv_packe
 use super::ics04_channel::handler::timeout::{
     timeout_packet_execute, timeout_packet_validate, TimeoutMsgType,
 };
-use super::ics04_channel::msgs::{ChannelMsg, PacketMsg};
+use super::ics04_channel::msgs::{
+    channel_msg_to_port_id, packet_msg_to_port_id, ChannelMsg, PacketMsg,
+};
 use super::msgs::MsgEnvelope;
 use super::router::Router;
-use super::{ContextError, ExecutionContext, ValidationContext};
+use super::{ExecutionContext, ValidationContext};
 
 /// Entrypoint which performs both validation and message execution
 pub fn dispatch(
@@ -78,12 +79,15 @@ where
         }
         .map_err(RouterError::ContextError),
         MsgEnvelope::Channel(msg) => {
+            let port_id = channel_msg_to_port_id(&msg);
             let module_id = router
-                .lookup_module_channel(&msg)
-                .map_err(ContextError::from)?;
+                .lookup_module(port_id)
+                .ok_or(RouterError::UnknownPort {
+                    port_id: port_id.clone(),
+                })?;
             let module = router
                 .get_route(&module_id)
-                .ok_or(ContextError::from(ChannelError::RouteNotFound))?;
+                .ok_or(RouterError::ModuleNotFound)?;
 
             match msg {
                 ChannelMsg::OpenInit(msg) => chan_open_init_validate(ctx, module, msg),
@@ -96,12 +100,15 @@ where
             .map_err(RouterError::ContextError)
         }
         MsgEnvelope::Packet(msg) => {
+            let port_id = packet_msg_to_port_id(&msg);
             let module_id = router
-                .lookup_module_packet(&msg)
-                .map_err(ContextError::from)?;
+                .lookup_module(port_id)
+                .ok_or(RouterError::UnknownPort {
+                    port_id: port_id.clone(),
+                })?;
             let module = router
                 .get_route(&module_id)
-                .ok_or(ContextError::from(ChannelError::RouteNotFound))?;
+                .ok_or(RouterError::ModuleNotFound)?;
 
             match msg {
                 PacketMsg::Recv(msg) => recv_packet_validate(ctx, msg),
@@ -147,12 +154,15 @@ where
         }
         .map_err(RouterError::ContextError),
         MsgEnvelope::Channel(msg) => {
+            let port_id = channel_msg_to_port_id(&msg);
             let module_id = router
-                .lookup_module_channel(&msg)
-                .map_err(ContextError::from)?;
+                .lookup_module(port_id)
+                .ok_or(RouterError::UnknownPort {
+                    port_id: port_id.clone(),
+                })?;
             let module = router
                 .get_route_mut(&module_id)
-                .ok_or(ContextError::from(ChannelError::RouteNotFound))?;
+                .ok_or(RouterError::ModuleNotFound)?;
 
             match msg {
                 ChannelMsg::OpenInit(msg) => chan_open_init_execute(ctx, module, msg),
@@ -165,12 +175,15 @@ where
             .map_err(RouterError::ContextError)
         }
         MsgEnvelope::Packet(msg) => {
+            let port_id = packet_msg_to_port_id(&msg);
             let module_id = router
-                .lookup_module_packet(&msg)
-                .map_err(ContextError::from)?;
+                .lookup_module(port_id)
+                .ok_or(RouterError::UnknownPort {
+                    port_id: port_id.clone(),
+                })?;
             let module = router
                 .get_route_mut(&module_id)
-                .ok_or(ContextError::from(ChannelError::RouteNotFound))?;
+                .ok_or(RouterError::ModuleNotFound)?;
 
             match msg {
                 PacketMsg::Recv(msg) => recv_packet_execute(ctx, module, msg),

crates/ibc/src/core/ics04_channel/error.rs

@@ -16,21 +16,8 @@ use crate::Height;
 
 use displaydoc::Display;
 
-#[derive(Debug, Display)]
-pub enum PortError {
-    /// port `{port_id}` is unknown
-    UnknownPort { port_id: PortId },
-    /// implementation specific error
-    ImplementationSpecific,
-}
-
-#[cfg(feature = "std")]
-impl std::error::Error for PortError {}
-
 #[derive(Debug, Display)]
 pub enum ChannelError {
-    /// port error: `{0}`
-    Port(PortError),
     /// invalid channel end: `{channel_end}`
     InvalidChannelEnd { channel_end: String },
     /// invalid channel id: expected `{expected}`, actual `{actual}`
@@ -79,8 +66,6 @@ pub enum ChannelError {
     ProcessedTimeNotFound { client_id: ClientId, height: Height },
     /// Processed height for the client `{client_id}` at height `{height}` not found
     ProcessedHeightNotFound { client_id: ClientId, height: Height },
-    /// route not found
-    RouteNotFound,
     /// application module error: `{description}`
     AppModule { description: String },
     /// other error: `{description}`
@@ -216,7 +201,6 @@ impl std::error::Error for PacketError {
 impl std::error::Error for ChannelError {
     fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
         match &self {
-            Self::Port(e) => Some(e),
             Self::InvalidIdentifier(e) => Some(e),
             Self::PacketVerificationFailed {
                 client_error: e, ..

crates/ibc/src/core/ics04_channel/msgs.rs

@@ -28,6 +28,8 @@ pub use recv_packet::MsgRecvPacket;
 pub use timeout::MsgTimeout;
 pub use timeout_on_close::MsgTimeoutOnClose;
 
+use crate::core::ics24_host::identifier::PortId;
+
 /// All channel messages
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum ChannelMsg {
@@ -47,3 +49,23 @@ pub enum PacketMsg {
     Timeout(MsgTimeout),
     TimeoutOnClose(MsgTimeoutOnClose),
 }
+
+pub(crate) fn channel_msg_to_port_id(msg: &ChannelMsg) -> &PortId {
+    match msg {
+        ChannelMsg::OpenInit(msg) => &msg.port_id_on_a,
+        ChannelMsg::OpenTry(msg) => &msg.port_id_on_b,
+        ChannelMsg::OpenAck(msg) => &msg.port_id_on_a,
+        ChannelMsg::OpenConfirm(msg) => &msg.port_id_on_b,
+        ChannelMsg::CloseInit(msg) => &msg.port_id_on_a,
+        ChannelMsg::CloseConfirm(msg) => &msg.port_id_on_b,
+    }
+}
+
+pub(crate) fn packet_msg_to_port_id(msg: &PacketMsg) -> &PortId {
+    match msg {
+        PacketMsg::Recv(msg) => &msg.packet.port_id_on_b,
+        PacketMsg::Ack(msg) => &msg.packet.port_id_on_a,
+        PacketMsg::Timeout(msg) => &msg.packet.port_id_on_a,
+        PacketMsg::TimeoutOnClose(msg) => &msg.packet.port_id_on_a,
+    }
+}

crates/ibc/src/core/router.rs

@@ -7,10 +7,7 @@ use core::fmt::{Debug, Display, Error as FmtError, Formatter};
 use crate::core::events::ModuleEvent;
 use crate::core::ics04_channel::acknowledgement::Acknowledgement;
 use crate::core::ics04_channel::channel::{Counterparty, Order};
-use crate::core::ics04_channel::error::PortError::UnknownPort;
 use crate::core::ics04_channel::error::{ChannelError, PacketError};
-use crate::core::ics04_channel::msgs::ChannelMsg;
-use crate::core::ics04_channel::msgs::PacketMsg;
 use crate::core::ics04_channel::packet::Packet;
 use crate::core::ics04_channel::Version;
 use crate::core::ics24_host::identifier::PortId;
@@ -26,39 +23,7 @@ pub trait Router {
     fn get_route_mut(&mut self, module_id: &ModuleId) -> Option<&mut dyn Module>;
 
     /// Return the module_id associated with a given port_id
-    fn lookup_module_by_port(&self, port_id: &PortId) -> Option<ModuleId>;
-
-    fn lookup_module_channel(&self, msg: &ChannelMsg) -> Result<ModuleId, ChannelError> {
-        let port_id = match msg {
-            ChannelMsg::OpenInit(msg) => &msg.port_id_on_a,
-            ChannelMsg::OpenTry(msg) => &msg.port_id_on_b,
-            ChannelMsg::OpenAck(msg) => &msg.port_id_on_a,
-            ChannelMsg::OpenConfirm(msg) => &msg.port_id_on_b,
-            ChannelMsg::CloseInit(msg) => &msg.port_id_on_a,
-            ChannelMsg::CloseConfirm(msg) => &msg.port_id_on_b,
-        };
-        let module_id = self
-            .lookup_module_by_port(port_id)
-            .ok_or(ChannelError::Port(UnknownPort {
-                port_id: port_id.clone(),
-            }))?;
-        Ok(module_id)
-    }
-
-    fn lookup_module_packet(&self, msg: &PacketMsg) -> Result<ModuleId, ChannelError> {
-        let port_id = match msg {
-            PacketMsg::Recv(msg) => &msg.packet.port_id_on_b,
-            PacketMsg::Ack(msg) => &msg.packet.port_id_on_a,
-            PacketMsg::Timeout(msg) => &msg.packet.port_id_on_a,
-            PacketMsg::TimeoutOnClose(msg) => &msg.packet.port_id_on_a,
-        };
-        let module_id = self
-            .lookup_module_by_port(port_id)
-            .ok_or(ChannelError::Port(UnknownPort {
-                port_id: port_id.clone(),
-            }))?;
-        Ok(module_id)
-    }
+    fn lookup_module(&self, port_id: &PortId) -> Option<ModuleId>;
 }
 
 /// Module name, internal to the chain.

crates/ibc/src/mock/router.rs

@@ -48,7 +48,7 @@ impl Router for MockRouter {
         }
     }
 
-    fn lookup_module_by_port(&self, port_id: &PortId) -> Option<ModuleId> {
+    fn lookup_module(&self, port_id: &PortId) -> Option<ModuleId> {
         self.port_to_module.get(port_id).cloned()
     }
 }