Implement forward-unix command on network client

Author: NickCao

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

@@ -1,5 +1,6 @@
 from contextlib import asynccontextmanager
 from functools import partial
+from os import PathLike
 
 from jumpstarter.client import DriverClient
 from jumpstarter.client.adapters import blocking
@@ -37,6 +38,7 @@ async def UnixPortforwardAdapter(
     *,
     client: DriverClient,
     method: str = "connect",
+    path: PathLike | None = None,
 ):
-    async with TemporaryUnixListener(partial(handler, client, method)) as addr:
+    async with TemporaryUnixListener(partial(handler, client, method), path=path) as addr:
         yield addr

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

@@ -4,7 +4,7 @@
 
 import asyncclick as click
 
-from .adapters import DbusAdapter, TcpPortforwardAdapter
+from .adapters import DbusAdapter, TcpPortforwardAdapter, UnixPortforwardAdapter
 from .driver import DbusNetwork
 from jumpstarter.client import DriverClient
 
@@ -20,7 +20,11 @@ def base():
         @click.option("--address", default="localhost", show_default=True)
         @click.argument("port", type=int)
         def forward_tcp(address: str, port: int):
-            """Forward local TCP port to remote network connection"""
+            """
+            Forward local TCP port to remote network
+
+            PORT is the TCP port to listen on.
+            """
 
             with TcpPortforwardAdapter(
                 client=self,
@@ -37,6 +41,24 @@ def forward_tcp(address: str, port: int):
 
                 Event().wait()
 
+        @base.command()
+        @click.argument("path", type=click.Path(), required=False)
+        def forward_unix(path: str | None):
+            """
+            Forward local Unix domain socket to remote network
+
+            PATH is the path of the Unix domain socket to listen on,
+            defaults to a random path under $XDG_RUNTIME_DIR.
+            """
+
+            with UnixPortforwardAdapter(
+                client=self,
+                path=path,
+            ) as addr:
+                click.echo(addr)
+
+                Event().wait()
+
         return base
 
 

packages/jumpstarter/jumpstarter/common/tempfile.py

@@ -1,4 +1,5 @@
-from contextlib import asynccontextmanager, contextmanager
+from contextlib import asynccontextmanager, contextmanager, nullcontext
+from os import PathLike
 from pathlib import Path
 from socket import AddressFamily
 from tempfile import TemporaryDirectory
@@ -15,8 +16,13 @@ def TemporarySocket():
 
 
 @asynccontextmanager
-async def TemporaryUnixListener(handler):
-    with TemporarySocket() as path:
+async def TemporaryUnixListener(handler, path: PathLike | None = None):
+    if path is not None:
+        cm = nullcontext(path)
+    else:
+        cm = TemporarySocket()
+
+    with cm as path:
         async with await create_unix_listener(path) as listener:
             async with create_task_group() as tg:
                 tg.start_soon(listener.serve, handler, tg)