Merge pull request #619 from jumpstarter-dev/backport-618-to-release-0.7

[Backport release-0.7] Add direct address reporting support for tcp/udp

Author: bennyz

packages/jumpstarter-driver-network/README.md

@@ -1,7 +1,8 @@
 # Network drivers
 
 `jumpstarter-driver-network` provides functionality for interacting with network
-servers and connections.
+servers and connections, redirecting DUT network services to the client handling
+the lease.
 
 ## Installation
 
@@ -19,9 +20,19 @@ export:
   network:
     type: jumpstarter_driver_network.driver.TcpNetwork
     config:
-      # Add required parameters here
+      host: 192.168.1.2
+      port: 5201
+      enable_address: true
 ```
 
+### Config parameters
+
+| Parameter     | Description                                         | Type  | Required | Default            |
+| ------------- | --------------------------------------------------- | ----- | -------- | ------------------ |
+| host          | Hostname or IP address of the DUT                   | str   | yes      |                    |
+| port          | Port number of the DUT service to connect to        | int   | yes      |                    |
+| enable_address | Whether to enable address mode (reporting the address of the client) | bool  | no | true   |
+
 ## API Reference
 
 Network driver classes:

packages/jumpstarter-driver-network/jumpstarter_driver_network/client.py

@@ -2,17 +2,23 @@
 from contextlib import contextmanager
 from ipaddress import IPv6Address, ip_address
 from threading import Event
-from typing import Any
+from typing import Any, Tuple
+from urllib.parse import urlparse
 
 import click
 from anyio import ContextManagerMixin
 
 from .adapters import DbusAdapter, TcpPortforwardAdapter, UnixPortforwardAdapter
 from .driver import DbusNetwork
 from jumpstarter.client import DriverClient
+from jumpstarter.client.core import DriverMethodNotImplemented
 
 
 class NetworkClient(DriverClient):
+
+    def address(self):
+        return self.call("address")
+
     def cli(self):
         @click.group
         def base():
@@ -62,9 +68,35 @@ def forward_unix(path: str | None):
 
                 Event().wait()
 
-        return base
+        @base.command()
+        @click.option("--host", is_flag=True)
+        @click.option("--port", is_flag=True)
+        def address(host, port):
+            """
+            Direct TCP connection to remote network
+            """
+            try:
+                addr = self.address()
+                if not host and not port or host and port:
+                    # Strip any URL scheme for clean display
+                    clean_addr = _strip_scheme(addr)
+                    click.echo(clean_addr)
+                else:
+                    # Parse address safely to handle IPv6
+                    parsed_host, parsed_port = _parse_address(addr)
+                    click.echo(parsed_host if host else parsed_port)
+            except ValueError as e:
+                raise click.ClickException(
+                    f"enable_address mode is not true in the exporter configuration: {e}"
+                ) from e
+            except DriverMethodNotImplemented as e:
+                raise click.ClickException(
+                    "This exporter does not support direct connection yet, update exporter to 0.7.1 or later"
+                ) from e
 
 
+        return base
+
 class DbusNetworkClient(NetworkClient, ContextManagerMixin):
     @contextmanager
     def __contextmanager__(self) -> Generator[Any]:
@@ -74,3 +106,57 @@ def __contextmanager__(self) -> Generator[Any]:
     @property
     def kind(self):
         return self.labels[DbusNetwork.KIND_LABEL]
+
+
+def _parse_address(addr: str) -> Tuple[str, str]:
+    """Parse a host:port address string, handling IPv6 addresses correctly.
+
+    Uses urllib.parse.urlparse for robust parsing of network addresses.
+
+    Returns:
+        Tuple of (host, port) as strings
+
+    Examples:
+        "127.0.0.1:8080" -> ("127.0.0.1", "8080")
+        "[::1]:8080" -> ("::1", "8080")
+        "localhost:8080" -> ("localhost", "8080")
+    """
+    # Add a dummy scheme to make it a valid URL for urlparse
+    if not addr.startswith(("http://", "https://", "tcp://", "udp://")):
+        addr = f"tcp://{addr}"
+
+    parsed = urlparse(addr)
+    host = parsed.hostname or ""
+    port = str(parsed.port) if parsed.port else ""
+
+    return host, port
+
+
+def _strip_scheme(addr: str) -> str:
+    """Remove URL scheme from address string for clean display.
+
+    Uses urllib.parse.urlparse to properly handle various URL formats.
+
+    Returns:
+        Address string without scheme prefix
+
+    Examples:
+        "tcp://127.0.0.1:8080" -> "127.0.0.1:8080"
+        "udp://[::1]:8080" -> "[::1]:8080"
+        "127.0.0.1:8080" -> "127.0.0.1:8080"
+    """
+    # Handle IPv6 addresses in brackets specially
+    if "://[" in addr and "]" in addr:
+        # Find the scheme separator and the closing bracket
+        scheme_end = addr.find("://")
+        if scheme_end != -1:
+            # Extract everything after "://"
+            return addr[scheme_end + 3:]
+
+    # For other cases, use urlparse
+    parsed = urlparse(addr)
+    # Reconstruct the address without scheme
+    if parsed.port:
+        return f"{parsed.hostname}:{parsed.port}"
+    else:
+        return parsed.hostname or addr

packages/jumpstarter-driver-network/jumpstarter_driver_network/driver.py

@@ -18,7 +18,60 @@
 from anyio.streams.stapled import StapledObjectStream
 
 from .streams.websocket import WebsocketClientStream
-from jumpstarter.driver import Driver, exportstream
+from jumpstarter.driver import Driver, export, exportstream
+
+
+def _is_ipv6_address(host: str) -> bool:
+    """Check if the given host string is an IPv6 address."""
+    try:
+        socket.inet_pton(socket.AF_INET6, host)
+        return True
+    except (OSError, ValueError):
+        return False
+
+
+def _is_ipv4_address(host: str) -> bool:
+    """Check if the given host string is an IPv4 address."""
+    try:
+        socket.inet_pton(socket.AF_INET, host)
+        return True
+    except (OSError, ValueError):
+        return False
+
+
+def _resolve_hostname(host: str) -> str:
+    """Resolve hostname to IP address. Returns the original host if resolution fails."""
+    # If it's already an IP address, return as-is
+    if _is_ipv4_address(host) or _is_ipv6_address(host):
+        return host
+
+    try:
+        # Try to resolve the hostname
+        # getaddrinfo returns a list of tuples, we want the first IP address
+        addr_info = socket.getaddrinfo(host, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
+        if addr_info:
+            # Get the first resolved address (ip address)
+            resolved_ip = addr_info[0][4][0]
+            return resolved_ip
+    except (OSError, socket.gaierror):
+        # If resolution fails, return the original hostname
+        pass
+
+    return host
+
+
+def _format_address(host: str, port: int) -> str:
+    """Format host and port as an address string, handling IPv6 addresses correctly.
+
+    Resolves hostnames to their IP addresses when possible.
+    """
+    # Resolve hostname to IP address
+    resolved_host = _resolve_hostname(host)
+
+    if _is_ipv6_address(resolved_host):
+        return f"[{resolved_host}]:{port}"
+    else:
+        return f"{resolved_host}:{port}"
 
 
 class NetworkInterface(metaclass=ABCMeta):
@@ -51,6 +104,7 @@ class TcpNetwork(NetworkInterface, Driver):
 
     host: str
     port: int
+    enable_address: bool = True
 
     @exportstream
     @asynccontextmanager
@@ -59,6 +113,12 @@ async def connect(self):
         async with await connect_tcp(remote_host=self.host, remote_port=self.port) as stream:
             yield stream
 
+    @export
+    async def address(self):
+        if self.enable_address:
+            return "tcp://" + _format_address(self.host, self.port)
+        else:
+            raise ValueError("enable_address mode is not true in the exporter configuration")
 
 @dataclass(kw_only=True)
 class UdpNetwork(NetworkInterface, Driver):
@@ -77,6 +137,7 @@ class UdpNetwork(NetworkInterface, Driver):
 
     host: str
     port: int
+    enable_address: bool = True
 
     @exportstream
     @asynccontextmanager
@@ -85,6 +146,13 @@ async def connect(self):
         async with await create_connected_udp_socket(remote_host=self.host, remote_port=self.port) as stream:
             yield stream
 
+    @export
+    async def address(self):
+        if self.enable_address:
+            return "udp://" + _format_address(self.host, self.port)
+        else:
+            raise ValueError("enable_address mode is not true in the exporter configuration")
+
 
 @dataclass(kw_only=True)
 class UnixNetwork(NetworkInterface, Driver):
@@ -245,12 +313,13 @@ class WebsocketNetwork(NetworkInterface, Driver):
     Handles websocket connections from a given url.
     '''
     url: str
+    enable_address: bool = True
 
     @exportstream
     @asynccontextmanager
     async def connect(self):
         '''
-        Create a websocket connection to `self.url` and srreams its output.
+        Create a websocket connection to `self.url` and streams its output.
         '''
         self.logger.info("Connecting to %s", self.url)
 
@@ -259,3 +328,10 @@ async def connect(self):
                 yield stream
 
         self.logger.info("Disconnected from %s", self.url)
+
+    @export
+    async def address(self):
+        if self.enable_address:
+            return self.url
+        else:
+            raise ValueError("enable_address mode is not true in the exporter configuration")