From 8224f1a05d36c08ec93fbab4564fa96d233bbfde Mon Sep 17 00:00:00 2001
From: Thomas Eizinger <thomas@eizinger.io>
Date: Mon, 8 May 2023 11:31:25 +0200
Subject: [PATCH] feat(identify): keep connection alive while we are using it

Currently, the `connection_keep_alive` function of identify does not compute anything but its return value is set through the handler state machine. This is hard to understand and causes surprising behaviour because at the moment, we set `KeepAlive::No` as soon as the remote has answered our identify request. Depending on what else is happening on the connection, this might close the connection before we have successfully answered the remote's identify request.

To fix this, we now compute `connection_keep_alive` based on whether we are still using the connection.

Related: #3844.

Pull-Request: #3876.
---
 protocols/identify/CHANGELOG.md   |  4 ++++
 protocols/identify/src/handler.rs | 20 +++++++++++++-------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/protocols/identify/CHANGELOG.md b/protocols/identify/CHANGELOG.md
index 69f26e44..e86afde8 100644
--- a/protocols/identify/CHANGELOG.md
+++ b/protocols/identify/CHANGELOG.md
@@ -7,9 +7,13 @@
 - Reduce the initial delay before running the identify protocol to 0 and make the option deprecated.
   See [PR 3545].
 
+- Fix aborting the answering of an identify request in rare situations.
+  See [PR 3876].
+
 [PR 3698]: https://github.com/libp2p/rust-libp2p/pull/3698
 [PR 3715]: https://github.com/libp2p/rust-libp2p/pull/3715
 [PR 3545]: https://github.com/libp2p/rust-libp2p/pull/3545
+[PR 3876]: https://github.com/libp2p/rust-libp2p/pull/3876
 
 ## 0.42.2
 
diff --git a/protocols/identify/src/handler.rs b/protocols/identify/src/handler.rs
index 0c72635d..96f15924 100644
--- a/protocols/identify/src/handler.rs
+++ b/protocols/identify/src/handler.rs
@@ -64,9 +64,6 @@ pub struct Handler {
     /// Future that fires when we need to identify the node again.
     trigger_next_identify: Delay,
 
-    /// Whether the handler should keep the connection alive.
-    keep_alive: KeepAlive,
-
     /// The interval of `trigger_next_identify`, i.e. the recurrent delay.
     interval: Duration,
 
@@ -132,7 +129,6 @@ impl Handler {
             reply_streams: VecDeque::new(),
             pending_replies: FuturesUnordered::new(),
             trigger_next_identify: Delay::new(initial_delay),
-            keep_alive: KeepAlive::Yes,
             interval,
             public_key,
             protocol_version,
@@ -190,7 +186,6 @@ impl Handler {
                     .push(ConnectionHandlerEvent::Custom(Event::Identified(
                         remote_info,
                     )));
-                self.keep_alive = KeepAlive::No;
             }
             future::Either::Right(()) => self
                 .events
@@ -210,7 +205,6 @@ impl Handler {
             .push(ConnectionHandlerEvent::Custom(Event::IdentificationError(
                 err,
             )));
-        self.keep_alive = KeepAlive::No;
         self.trigger_next_identify.reset(self.interval);
     }
 }
@@ -268,7 +262,19 @@ impl ConnectionHandler for Handler {
     }
 
     fn connection_keep_alive(&self) -> KeepAlive {
-        self.keep_alive
+        if self.inbound_identify_push.is_some() {
+            return KeepAlive::Yes;
+        }
+
+        if !self.pending_replies.is_empty() {
+            return KeepAlive::Yes;
+        }
+
+        if !self.reply_streams.is_empty() {
+            return KeepAlive::Yes;
+        }
+
+        KeepAlive::No
     }
 
     fn poll(