From 5b589ba8f5bb5c258a37e65f73267ea24eaff303 Mon Sep 17 00:00:00 2001
From: Pete Burgers <pete@anvil.works>
Date: Thu, 17 Oct 2024 19:31:23 +0100
Subject: [PATCH] Allow binding to local port 0

Add a from_port property that contains the actual local port used for
the portforward.
---
 python/portforward/__init__.py | 12 +++++++++---
 src/portforward.rs             |  8 +++++---
 tests/test_portforward.py      | 27 +++++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/python/portforward/__init__.py b/python/portforward/__init__.py
index 17cac6e..8f37cbf 100644
--- a/python/portforward/__init__.py
+++ b/python/portforward/__init__.py
@@ -57,7 +57,7 @@ def forward(
 
     :param namespace: Target namespace
     :param pod_or_service: Name of target Pod or service
-    :param from_port: Local port
+    :param from_port: Local port, or 0 to use any free port
     :param to_port: Port inside the pod
     :param config_path: Path for loading kube config
     :param waiting: Delay in seconds
@@ -129,6 +129,11 @@ def stop(self):
     def is_stopped(self):
         return self._async_forwarder.is_stopped
 
+    @property
+    def from_port(self):
+        """The local port that was actually used for the portforward."""
+        return self._async_forwarder.from_port
+
 
 class AsyncPortForwarder:
     """Use the same args as the `portforward.forward` method."""
@@ -158,11 +163,12 @@ def __init__(
         bind_ip = _validate_ip_address(bind_ip)
 
         self.actual_pod_name: str = ""
+        self.from_port: int = 0
         self._is_stopped: bool = False
         self.bind_address: str = f"{bind_ip}:{from_port}"
 
     async def forward(self):
-        self.actual_pod_name = await _portforward.forward(
+        (self.actual_pod_name, self.from_port) = await _portforward.forward(
             self.namespace,
             self.pod_or_service,
             self.bind_address,
@@ -200,7 +206,7 @@ def _validate_str(arg_name, arg) -> str:
 
 
 def _validate_port(arg_name, arg) -> int:
-    in_range = arg and 0 < arg < 65536
+    in_range = arg is not None and 0 <= arg < 65536
     if arg is None or not isinstance(arg, int) or not in_range:
         raise ValueError(f"{arg_name}={arg} is not a valid port")
 
diff --git a/src/portforward.rs b/src/portforward.rs
index aac8377..b2bbcd1 100644
--- a/src/portforward.rs
+++ b/src/portforward.rs
@@ -34,9 +34,10 @@ pub struct ForwardConfig {
     kube_context: String,
 }
 
-/// Creates a connection to a pod. It returns the actual pod name for the portforward.
+/// Creates a connection to a pod. It returns a `(pod_name, from_port)` tuple
+/// with the actual pod name and local port used for the portforward.
 /// It differs from `pod_or_service` when `pod_or_service` represents a service.
-pub async fn forward(config: ForwardConfig) -> anyhow::Result<String> {
+pub async fn forward(config: ForwardConfig) -> anyhow::Result<(String, u16)> {
     debug!("{:?}", config);
 
     let client_config = load_config(&config.config_path, &config.kube_context).await?;
@@ -58,6 +59,7 @@ pub async fn forward(config: ForwardConfig) -> anyhow::Result<String> {
 
     let addr = SocketAddr::from_str(&config.bind_address).with_context(move || config.bind_address)?;
     let tcp_listener = TcpListener::bind(addr).await?;
+    let from_port = tcp_listener.local_addr()?.port();
     let forward_task = setup_forward_task(
         tcp_listener,
         rx,
@@ -68,7 +70,7 @@ pub async fn forward(config: ForwardConfig) -> anyhow::Result<String> {
 
     tokio::spawn(forward_task);
 
-    return Ok(q_name.pod_name);
+    return Ok((q_name.pod_name, from_port));
 }
 
 async fn load_config(
diff --git a/tests/test_portforward.py b/tests/test_portforward.py
index 95f81c0..174d515 100644
--- a/tests/test_portforward.py
+++ b/tests/test_portforward.py
@@ -124,6 +124,33 @@ def test_service_portforward_with_success(kind_cluster: KindCluster):
         response: requests.Response = requests.get(url_2)
         pytest.fail("Portforward should be closed after leaving the context manager")
 
+def test_portforward_from_port_zero_assigns_port(kind_cluster: KindCluster):
+    # Arrange
+    _create_test_resources(kind_cluster)
+
+    pod_name = "test-pod"
+    config = str(kind_cluster.kubeconfig_path.absolute())
+
+    local_port = 0  # from port
+    pod_port = 3000  # to port
+
+    pf = portforward.forward(
+        TEST_NAMESPACE,
+        pod_name,
+        local_port,
+        pod_port,
+        config_path=config,
+        kube_context=TEST_CONTEXT,
+    )
+
+    # Act & Assert
+    with pf as forwarder:
+        assert not forwarder.is_stopped()
+        assert forwarder.from_port != 0
+        url = f"http://localhost:{forwarder.from_port}/ping"
+        response: requests.Response = requests.get(url)
+        assert response.status_code == 200
+
 
 @pytest.mark.parametrize(
     "namespace,pod,from_port,to_port",