diff --git a/core/src/connection.rs b/core/src/connection.rs
index 513665a631a..9e39ae21807 100644
--- a/core/src/connection.rs
+++ b/core/src/connection.rs
@@ -229,10 +229,8 @@ where
         self.handler.inject_event(event);
     }
 
-    // TODO: Update comment
-    //
-    /// Begins an orderly shutdown of the connection, returning a
-    /// `Future` that resolves when connection shutdown is complete.
+    /// Begins an orderly shutdown of the connection, returning the connection
+    /// handler and a `Future` that resolves when connection shutdown is complete.
     pub fn close(self) -> (THandler, Close<TMuxer>) {
         (self.handler, self.muxing.close().0)
     }
diff --git a/core/src/connection/manager/task.rs b/core/src/connection/manager/task.rs
index f81c520e6e9..1dc56df81ca 100644
--- a/core/src/connection/manager/task.rs
+++ b/core/src/connection/manager/task.rs
@@ -329,8 +329,7 @@ where
                             Poll::Ready(Err(error)) => {
                                 // Don't accept any further commands.
                                 this.commands.get_mut().close();
-                                // TODO: Good idea if there is already an error?
-                                let (handler, _) = connection.close();
+                                let (handler, _closing_muxer) = connection.close();
                                 // Terminate the task with the error, dropping the connection.
                                 let event = Event::Closed {
                                     id,
diff --git a/swarm/src/behaviour.rs b/swarm/src/behaviour.rs
index ac99bd3e701..3dd02bac308 100644
--- a/swarm/src/behaviour.rs
+++ b/swarm/src/behaviour.rs
@@ -26,6 +26,10 @@ use libp2p_core::{
 };
 use std::{error, task::Context, task::Poll};
 
+/// Custom event that can be received by the [`ProtocolsHandler`].
+type THandlerInEvent<THandler> =
+    <<THandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::InEvent;
+
 /// A behaviour for the network. Allows customizing the swarm.
 ///
 /// This trait has been designed to be composable. Multiple implementations can be combined into
@@ -112,7 +116,14 @@ pub trait NetworkBehaviour: Send + 'static {
     /// A call to this method is always paired with an earlier call to
     /// `inject_connection_established` with the same peer ID, connection ID and
     /// endpoint.
-    fn inject_connection_closed(&mut self, _: &PeerId, _: &ConnectionId, _: &ConnectedPoint, _: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler) {}
+    fn inject_connection_closed(
+        &mut self,
+        _: &PeerId,
+        _: &ConnectionId,
+        _: &ConnectedPoint,
+        _: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
+    ) {
+    }
 
     /// Informs the behaviour that the [`ConnectedPoint`] of an existing connection has changed.
     fn inject_address_change(
@@ -182,8 +193,11 @@ pub trait NetworkBehaviour: Send + 'static {
     ///
     /// This API mimics the API of the `Stream` trait. The method may register the current task in
     /// order to wake it up at a later point in time.
-    fn poll(&mut self, cx: &mut Context<'_>, params: &mut impl PollParameters)
-        -> Poll<NetworkBehaviourAction<<<Self::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::InEvent, Self::OutEvent, Self::ProtocolsHandler>>;
+    fn poll(
+        &mut self,
+        cx: &mut Context<'_>,
+        params: &mut impl PollParameters,
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>>;
 }
 
 /// Parameters passed to `poll()`, that the `NetworkBehaviour` has access to.
@@ -228,9 +242,13 @@ pub trait NetworkBehaviourEventProcess<TEvent> {
 /// in whose context it is executing.
 ///
 /// [`Swarm`]: super::Swarm
+//
+// Note: `TInEvent` is needed to be able to implement
+// [`NetworkBehaviourAction::map_in`], mapping the handler `InEvent` leaving the
+// handler itself untouched.
 #[derive(Debug)]
-// TODO: Derive TInEvent from THandler.
-pub enum NetworkBehaviourAction<TInEvent, TOutEvent, THandler> {
+pub enum NetworkBehaviourAction<TOutEvent, THandler: IntoProtocolsHandler, TInEvent = THandlerInEvent<THandler>>
+{
     /// Instructs the `Swarm` to return an event when it is being polled.
     GenerateEvent(TOutEvent),
 
@@ -317,17 +335,29 @@ pub enum NetworkBehaviourAction<TInEvent, TOutEvent, THandler> {
     },
 }
 
-impl<TInEvent, TOutEvent, THandler> NetworkBehaviourAction<TInEvent, TOutEvent, THandler> {
+impl<TOutEvent, THandler: IntoProtocolsHandler> NetworkBehaviourAction<TOutEvent, THandler> {
     /// Map the handler event.
-    pub fn map_in<E>(self, f: impl FnOnce(TInEvent) -> E) -> NetworkBehaviourAction<E, TOutEvent, THandler> {
+    pub fn map_in<New>(
+        self,
+        f: impl FnOnce(THandlerInEvent<THandler>) -> New,
+    ) -> NetworkBehaviourAction<TOutEvent, THandler, New>
+    where
+        Self: ,
+    {
         match self {
             NetworkBehaviourAction::GenerateEvent(e) => NetworkBehaviourAction::GenerateEvent(e),
             NetworkBehaviourAction::DialAddress { address, handler } => {
                 NetworkBehaviourAction::DialAddress { address, handler }
             }
-            NetworkBehaviourAction::DialPeer { peer_id, condition, handler } => {
-                NetworkBehaviourAction::DialPeer { peer_id, condition, handler }
-            }
+            NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler,
+            } => NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler,
+            },
             NetworkBehaviourAction::NotifyHandler {
                 peer_id,
                 handler,
@@ -351,15 +381,21 @@ impl<TInEvent, TOutEvent, THandler> NetworkBehaviourAction<TInEvent, TOutEvent,
     }
 
     /// Map the event the swarm will return.
-    pub fn map_out<E>(self, f: impl FnOnce(TOutEvent) -> E) -> NetworkBehaviourAction<TInEvent, E, THandler> {
+    pub fn map_out<E>(self, f: impl FnOnce(TOutEvent) -> E) -> NetworkBehaviourAction<E, THandler> {
         match self {
             NetworkBehaviourAction::GenerateEvent(e) => NetworkBehaviourAction::GenerateEvent(f(e)),
             NetworkBehaviourAction::DialAddress { address, handler } => {
                 NetworkBehaviourAction::DialAddress { address, handler }
             }
-            NetworkBehaviourAction::DialPeer { peer_id, condition, handler } => {
-                NetworkBehaviourAction::DialPeer { peer_id, condition, handler }
-            }
+            NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler,
+            } => NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler,
+            },
             NetworkBehaviourAction::NotifyHandler {
                 peer_id,
                 handler,
diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
index 47cac9fbb1d..b93ae56ff4b 100644
--- a/swarm/src/lib.rs
+++ b/swarm/src/lib.rs
@@ -1275,7 +1275,6 @@ impl NetworkBehaviour for DummyBehaviour {
         _: &mut impl PollParameters,
     ) -> Poll<
         NetworkBehaviourAction<
-            <Self::ProtocolsHandler as ProtocolsHandler>::InEvent,
             Self::OutEvent,
             Self::ProtocolsHandler,
         >,
diff --git a/swarm/src/toggle.rs b/swarm/src/toggle.rs
index 6e72f2c6b90..84cdbfd3163 100644
--- a/swarm/src/toggle.rs
+++ b/swarm/src/toggle.rs
@@ -204,9 +204,11 @@ where
         }
     }
 
-    fn poll(&mut self, cx: &mut Context<'_>, params: &mut impl PollParameters)
-        -> Poll<NetworkBehaviourAction<<<Self::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::InEvent, Self::OutEvent, Self::ProtocolsHandler>>
-    {
+    fn poll(
+        &mut self,
+        cx: &mut Context<'_>,
+        params: &mut impl PollParameters,
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
         if let Some(inner) = self.inner.as_mut() {
             // inner.poll(cx, params)
             todo!("Wrap handler in Dial events with ToggleHandler.");