From fe3bcc9cd7eeb4a81b4ea8e3097e251830119ab7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Smolarek?=
 <34063647+Razz4780@users.noreply.github.com>
Date: Mon, 3 Feb 2025 19:05:57 +0100
Subject: [PATCH] Added option to set agent port (#3053)

---
 changelog.d/+set-agent-port.internal.md |  1 +
 mirrord/cli/src/connection.rs           |  2 +-
 mirrord/kube/src/api/container.rs       |  3 ++-
 mirrord/kube/src/api/kubernetes.rs      | 10 ++++++++--
 4 files changed, 12 insertions(+), 4 deletions(-)
 create mode 100644 changelog.d/+set-agent-port.internal.md

diff --git a/changelog.d/+set-agent-port.internal.md b/changelog.d/+set-agent-port.internal.md
new file mode 100644
index 00000000000..6fc1f928401
--- /dev/null
+++ b/changelog.d/+set-agent-port.internal.md
@@ -0,0 +1 @@
+`mirrord-kube` now allows for setting agent listen port.
diff --git a/mirrord/cli/src/connection.rs b/mirrord/cli/src/connection.rs
index dacec4fcb1b..fed2cc1b384 100644
--- a/mirrord/cli/src/connection.rs
+++ b/mirrord/cli/src/connection.rs
@@ -173,7 +173,7 @@ where
 
     let agent_connect_info = tokio::time::timeout(
         Duration::from_secs(config.agent.startup_timeout),
-        k8s_api.create_agent(progress, &config.target, Some(config), Default::default()),
+        k8s_api.create_agent(progress, &config.target, Some(config), None, None),
     )
     .await
     .unwrap_or(Err(KubeApiError::AgentReadyTimeout))
diff --git a/mirrord/kube/src/api/container.rs b/mirrord/kube/src/api/container.rs
index a651dc13458..49fba30e18b 100644
--- a/mirrord/kube/src/api/container.rs
+++ b/mirrord/kube/src/api/container.rs
@@ -53,8 +53,9 @@ impl ContainerParams {
         tls_cert: Option<String>,
         pod_ips: Option<String>,
         support_ipv6: bool,
+        port: Option<u16>,
     ) -> ContainerParams {
-        let port: u16 = rand::thread_rng().gen_range(30000..=65535);
+        let port = port.unwrap_or_else(|| rand::thread_rng().gen_range(30000..=65535));
         let gid: u16 = rand::thread_rng().gen_range(3000..u16::MAX);
 
         let name = format!(
diff --git a/mirrord/kube/src/api/kubernetes.rs b/mirrord/kube/src/api/kubernetes.rs
index 44410f34b74..86a84f7e4ba 100644
--- a/mirrord/kube/src/api/kubernetes.rs
+++ b/mirrord/kube/src/api/kubernetes.rs
@@ -161,12 +161,15 @@ impl KubernetesAPI {
     /// * `tls_cert` - value for
     ///   [`AGENT_OPERATOR_CERT_ENV`](mirrord_protocol::AGENT_OPERATOR_CERT_ENV), for creating an
     ///   agent from the operator. In usage from this repo this is always `None`.
+    /// * `agent_port` - port number on which the agent will listen for client connections. If
+    ///   [`None`] is given, a random high port will be user.
     #[tracing::instrument(level = "trace", skip(self), ret, err)]
     pub async fn create_agent_params(
         &self,
         target: &TargetConfig,
         tls_cert: Option<String>,
         support_ipv6: bool,
+        agent_port: Option<u16>,
     ) -> Result<(ContainerParams, Option<RuntimeData>), KubeApiError> {
         let runtime_data = match target.path.as_ref().unwrap_or(&Target::Targetless) {
             Target::Targetless => None,
@@ -188,7 +191,7 @@ impl KubernetesAPI {
                     .join(",")
             });
 
-        let params = ContainerParams::new(tls_cert, pod_ips, support_ipv6);
+        let params = ContainerParams::new(tls_cert, pod_ips, support_ipv6, agent_port);
 
         Ok((params, runtime_data))
     }
@@ -199,6 +202,8 @@ impl KubernetesAPI {
     /// * `tls_cert` - value for
     ///   [`AGENT_OPERATOR_CERT_ENV`](mirrord_protocol::AGENT_OPERATOR_CERT_ENV), for creating an
     ///   agent from the operator. In usage from this repo this is always `None`.
+    /// * `agent_port` - port number on which the agent will listen for client connections. If
+    ///   [`None`] is given, a random high port will be used.
     #[tracing::instrument(level = "trace", skip(self, progress))]
     pub async fn create_agent<P>(
         &self,
@@ -206,6 +211,7 @@ impl KubernetesAPI {
         target: &TargetConfig,
         config: Option<&LayerConfig>,
         tls_cert: Option<String>,
+        agent_port: Option<u16>,
     ) -> Result<AgentKubernetesConnectInfo, KubeApiError>
     where
         P: Progress + Send + Sync,
@@ -214,7 +220,7 @@ impl KubernetesAPI {
             .map(|layer_conf| layer_conf.feature.network.ipv6)
             .unwrap_or_default();
         let (params, runtime_data) = self
-            .create_agent_params(target, tls_cert, support_ipv6)
+            .create_agent_params(target, tls_cert, support_ipv6, agent_port)
             .await?;
         if let Some(RuntimeData {
             guessed_container: true,