From c7efd30ef88ac868e66023cfae46c7238a5ff4f8 Mon Sep 17 00:00:00 2001
From: Gemma <58080601+gememma@users.noreply.github.com>
Date: Fri, 6 Dec 2024 12:39:12 +0000
Subject: [PATCH] Stop mirrord ending unexpectedly when port forwarding in
 reverse (#2964)

* Fix ping/pong routine in ReversePortForwarder

* Add changelog
---
 changelog.d/2962.fixed.md       | 1 +
 mirrord/cli/src/port_forward.rs | 5 ++++-
 2 files changed, 5 insertions(+), 1 deletion(-)
 create mode 100644 changelog.d/2962.fixed.md

diff --git a/changelog.d/2962.fixed.md b/changelog.d/2962.fixed.md
new file mode 100644
index 00000000000..7db8d006ace
--- /dev/null
+++ b/changelog.d/2962.fixed.md
@@ -0,0 +1 @@
+Stop mirrord ending unexpectedly when port forwarding in reverse due to unexpected connection end.
\ No newline at end of file
diff --git a/mirrord/cli/src/port_forward.rs b/mirrord/cli/src/port_forward.rs
index c014fab7c8f..ba59679d945 100644
--- a/mirrord/cli/src/port_forward.rs
+++ b/mirrord/cli/src/port_forward.rs
@@ -440,6 +440,7 @@ pub struct ReversePortForwarder {
 
     /// true if Ping has been sent to agent.
     waiting_for_pong: bool,
+    ping_pong_timeout: Instant,
 }
 
 impl ReversePortForwarder {
@@ -480,6 +481,7 @@ impl ReversePortForwarder {
             background_tasks,
             incoming_proxy: incoming,
             waiting_for_pong: false,
+            ping_pong_timeout: Instant::now(),
         })
     }
 
@@ -511,13 +513,14 @@ impl ReversePortForwarder {
 
         loop {
             select! {
-                _ = tokio::time::sleep(Duration::from_secs(30)) => {
+                _ = tokio::time::sleep_until(self.ping_pong_timeout.into()) => {
                     if self.waiting_for_pong {
                         // no pong received before timeout
                         break Err(PortForwardError::AgentError("agent failed to respond to Ping".into()));
                     }
                     self.agent_connection.sender.send(ClientMessage::Ping).await?;
                     self.waiting_for_pong = true;
+                    self.ping_pong_timeout = Instant::now() + Duration::from_secs(30);
                 },
 
                 message = self.agent_connection.receiver.recv() => match message {