Send SwitchProtocolVersion on ConnectionRefresh

metalbear-co · Jan 28, 2025 · 1e729b8 · 1e729b8
1 parent 377fe96
commit 1e729b8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/mirrord/intproxy/src/lib.rs b/mirrord/intproxy/src/lib.rs
@@ -412,6 +412,11 @@ impl IntProxy {
 
     #[tracing::instrument(level = Level::TRACE, skip(self), err)]
     async fn handle_connection_refresh(&self) -> Result<(), IntProxyError> {
+        self.task_txs
+            .simple
+            .send(SimpleProxyMessage::ConnectionRefresh)
+            .await;
+
         self.task_txs
             .incoming
             .send(IncomingProxyMessage::ConnectionRefresh)

diff --git a/mirrord/intproxy/src/proxies/simple.rs b/mirrord/intproxy/src/proxies/simple.rs
@@ -27,6 +27,8 @@ pub enum SimpleProxyMessage {
     GetEnvRes(RemoteResult<HashMap<String, String>>),
     /// Protocol version was negotiated with the agent.
     ProtocolVersion(Version),
+    /// Agent connection was refreshed need to negotiate version
+    ConnectionRefresh,
 }
 
 #[derive(Error, Debug)]
@@ -123,6 +125,13 @@ impl BackgroundTask for SimpleProxy {
                         .await
                 }
                 SimpleProxyMessage::ProtocolVersion(version) => self.set_protocol_version(version),
+                SimpleProxyMessage::ConnectionRefresh => {
+                    if let Some(version) = &self.protocol_version {
+                        message_bus
+                            .send(ClientMessage::SwitchProtocolVersion(version.clone()))
+                            .await
+                    }
+                }
             }
         }