swarm/src/toggle: Fix TODOs

mxinden · Aug 20, 2021 · aa02e5f · aa02e5f
1 parent b73139a
commit aa02e5f
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 6 deletions.
diff --git a/swarm/src/behaviour.rs b/swarm/src/behaviour.rs
@@ -423,6 +423,60 @@ impl<TOutEvent, THandler: IntoProtocolsHandler> NetworkBehaviourAction<TOutEvent
     }
 }
 
+impl<TInEvent, TOutEvent, THandlerOld> NetworkBehaviourAction<TOutEvent, THandlerOld>
+where
+    THandlerOld: IntoProtocolsHandler,
+    <THandlerOld as IntoProtocolsHandler>::Handler: ProtocolsHandler<InEvent = TInEvent>,
+{
+    /// Map the handler event.
+    pub fn map_handler<THandlerNew>(
+        self,
+        f: impl FnOnce(THandlerOld) -> THandlerNew,
+    ) -> NetworkBehaviourAction<TOutEvent, THandlerNew>
+    where
+        THandlerNew: IntoProtocolsHandler,
+        <THandlerNew as IntoProtocolsHandler>::Handler: ProtocolsHandler<InEvent = TInEvent>,
+    {
+        match self {
+            NetworkBehaviourAction::GenerateEvent(e) => NetworkBehaviourAction::GenerateEvent(e),
+            NetworkBehaviourAction::DialAddress { address, handler } => {
+                NetworkBehaviourAction::DialAddress {
+                    address,
+                    handler: f(handler),
+                }
+            }
+            NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler,
+            } => NetworkBehaviourAction::DialPeer {
+                peer_id,
+                condition,
+                handler: f(handler),
+            },
+            NetworkBehaviourAction::NotifyHandler {
+                peer_id,
+                handler,
+                event,
+            } => NetworkBehaviourAction::NotifyHandler {
+                peer_id,
+                handler,
+                event: event,
+            },
+            NetworkBehaviourAction::ReportObservedAddr { address, score } => {
+                NetworkBehaviourAction::ReportObservedAddr { address, score }
+            }
+            NetworkBehaviourAction::CloseConnection {
+                peer_id,
+                connection,
+            } => NetworkBehaviourAction::CloseConnection {
+                peer_id,
+                connection,
+            },
+        }
+    }
+}
+
 /// The options w.r.t. which connection handler to notify of an event.
 #[derive(Debug, Clone)]
 pub enum NotifyHandler {

diff --git a/swarm/src/toggle.rs b/swarm/src/toggle.rs
@@ -116,8 +116,9 @@ where
         handler: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
     ) {
         if let Some(inner) = self.inner.as_mut() {
-            // TODO: Unwrap safe here?
-            inner.inject_connection_closed(peer_id, connection, endpoint, handler.inner.unwrap())
+            if let Some(handler) = handler.inner {
+                inner.inject_connection_closed(peer_id, connection, endpoint, handler)
+            }
         }
     }
 
@@ -157,8 +158,9 @@ where
 
     fn inject_dial_failure(&mut self, peer_id: &PeerId, handler: Self::ProtocolsHandler) {
         if let Some(inner) = self.inner.as_mut() {
-            // TODO: Unwrap safe here?
-            inner.inject_dial_failure(peer_id, handler.inner.unwrap())
+            if let Some(handler) = handler.inner {
+                inner.inject_dial_failure(peer_id, handler)
+            }
         }
     }
 
@@ -210,8 +212,9 @@ where
         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.");
+            inner
+                .poll(cx, params)
+                .map(|action| action.map_handler(|h| ToggleIntoProtoHandler { inner: Some(h) }))
         } else {
             Poll::Pending
         }