From 1f74f96ab4d2fae0d39dc9c56e0aec01e98462a9 Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Thu, 2 Nov 2023 17:16:19 +0100
Subject: [PATCH] fix(swarm): implement ConnectionHandler::poll_close for
 combinators

Follow-up to https://github.com/libp2p/rust-libp2p/pull/4076.

This is especially relevant since `libp2p-swarm-derive` uses `SelectConnectionHandler`.
---
 swarm/src/behaviour/toggle.rs |  8 ++++++++
 swarm/src/handler/either.rs   | 14 ++++++++++++++
 swarm/src/handler/map_in.rs   |  7 +++++++
 swarm/src/handler/map_out.rs  |  6 ++++++
 4 files changed, 35 insertions(+)

diff --git a/swarm/src/behaviour/toggle.rs b/swarm/src/behaviour/toggle.rs
index e1da71a0450..d3333acedaa 100644
--- a/swarm/src/behaviour/toggle.rs
+++ b/swarm/src/behaviour/toggle.rs
@@ -369,4 +369,12 @@ where
             }
         }
     }
+
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        if let Some(inner) = self.inner.as_mut() {
+            inner.poll_close(cx)
+        } else {
+            Poll::Ready(None)
+        }
+    }
 }
diff --git a/swarm/src/handler/either.rs b/swarm/src/handler/either.rs
index 093900135b8..41c5ff93369 100644
--- a/swarm/src/handler/either.rs
+++ b/swarm/src/handler/either.rs
@@ -142,6 +142,20 @@ where
         Poll::Ready(event)
     }
 
+    fn poll_close(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::ToBehaviour>> {
+        let event = match self {
+            Either::Left(handler) => futures::ready!(handler.poll_close(cx))
+                .map(Either::Left),
+            Either::Right(handler) => futures::ready!(handler.poll_close(cx))
+                .map(Either::Right),
+        };
+
+        Poll::Ready(event)
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<
diff --git a/swarm/src/handler/map_in.rs b/swarm/src/handler/map_in.rs
index e3458eb5451..1b5dc69a317 100644
--- a/swarm/src/handler/map_in.rs
+++ b/swarm/src/handler/map_in.rs
@@ -86,6 +86,13 @@ where
         self.inner.poll(cx)
     }
 
+    fn poll_close(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::ToBehaviour>> {
+        self.inner.poll_close(cx)
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<
diff --git a/swarm/src/handler/map_out.rs b/swarm/src/handler/map_out.rs
index cc06a4c50c8..56673c90372 100644
--- a/swarm/src/handler/map_out.rs
+++ b/swarm/src/handler/map_out.rs
@@ -90,6 +90,12 @@ where
         })
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        self.inner
+            .poll_close(cx)
+            .map(|ev| ev.map(|ev| (self.map)(ev)))
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<