proxy

Flight SSH proxy functionality.

FlightProxy

FlightProxy(url, backend_url, proxy_map, **kwargs)

Bases: FlightServerBase

Transparent Flight proxy that rewrites endpoint URIs based on mapping.

Rewrites FlightInfo endpoint location URIs on GetFlightInfo responses; all other RPC methods are forwarded to the backend unchanged. This is useful for proxying a remote Flight server through SSH port forwards.

Parameters:

Name	Type	Description	Default
url	str	URL to listen for queries.	required
backend_url	str	URL of remote info server to proxy.	required
proxy_map	dict[str, str]	Dictionary mapping of endpoint locations to be replaced. It should be keyed on OLD_HOST:PORT with value as NEW_HOST:PORT. The URL scheme will be preserved.	required

Source code in arrakis/proxy.py

def __init__(
    self,
    url: str,
    backend_url: str,
    proxy_map: dict[str, str],
    **kwargs,
):
    super().__init__(url, **kwargs)
    self.proxy_map = proxy_map
    self.client = flight.connect(backend_url)
    logger.debug("flight proxy server initialized: %s -> %s", url, backend_url)
    logger.debug("proxy_map: %s", self.proxy_map)

SSHConnection

SSHConnection(destination)

Manage a background SSH connection

Parameters:

Name	Type	Description	Default
destination	str	SSH destination (see ssh(1) for form).	required

Source code in arrakis/proxy.py

def __init__(self, destination: str):
    self.destination = destination
    self.ctrl_dir = pathlib.Path(tempfile.mkdtemp())
    self.ctrl_path = self.ctrl_dir / urlparse(self.destination).netloc
    self.ctrl = str(self.ctrl_path)
    self._forward_map: dict[str, str] = {}

forward_map property

forward_map

Dictionary of port forwards

REMOTE_HOSTPORT: LOCAL_HOSTPORT

close

close()

Close the connection

Source code in arrakis/proxy.py

def close(self):
    """Close the connection"""
    self.exec(["-O", "exit"], check=False, capture_output=True)
    self.ctrl_dir.rmdir()

connect

connect()

connect to the ssh destination

Source code in arrakis/proxy.py

def connect(self):
    """connect to the ssh destination"""
    cmd = [
        "ssh",
        "-S",
        self.ctrl,
        "-M",
        "-o",
        "ControlPersist=yes",
        "-f",
        # "-N",
        self.destination,
        "sleep",
        "60",
    ]
    logger.debug(" ".join(cmd))
    subprocess.run(cmd, check=True)  # noqa S603

exec

exec(ssh_cmd=None, shell_cmd=None, **kwargs)

exec ssh command on control master

Source code in arrakis/proxy.py

def exec(
    self,
    ssh_cmd: list[str] | None = None,
    shell_cmd: list[str] | None = None,
    **kwargs,
):
    """exec ssh command on control master"""
    cmd = [
        "ssh",
        "-S",
        self.ctrl,
        "-o",
        "ControlMaster=no",
    ]
    if ssh_cmd:
        cmd += ssh_cmd
    cmd += [self.destination]
    if shell_cmd:
        cmd += shell_cmd
    logger.debug(" ".join(cmd))
    return subprocess.run(  # noqa S603
        cmd,
        **kwargs,
    )

forward_port

forward_port(remote_hostport, local_port=None, *, wait=False)

initiate a local port forward to remote location

Source code in arrakis/proxy.py

def forward_port(
    self, remote_hostport: str, local_port: int | None = None, *, wait: bool = False
):
    """initiate a local port forward to remote location"""
    if not local_port:
        # hacky way to find an unused port. the acquired port
        # could be stolen between when the socket releases it and
        # ssh tries to take it.
        s = socket.socket()
        s.bind(("localhost", 0))
        _, local_port = s.getsockname()
        s.close()

    local_hostport = f"localhost:{local_port}"

    forward = f"{local_hostport}:{remote_hostport}"
    logger.debug("forwarding port: %s", forward)
    self.exec(["-O", "forward", "-L", forward], check=True)

    if wait:
        # wait for forward to be established
        s = socket.socket()
        while True:
            try:
                s.connect(("localhost", local_port))
                s.close()
                break
            except ConnectionRefusedError:
                time.sleep(0.01)
                continue

    self.forward_map[remote_hostport] = local_hostport
    return local_hostport

ssh_proxy

ssh_proxy(ssh_dest, arrakis_server=None, bind_address=None)

Create Flight proxy server over SSH.

This is done by:

1. Make SSH connection to the remote host.
2. Determine initial Flight info server URL on the remote side.
3. Retrieve all known endpoints from the remote info server.
4. Set up local ssh port forwards to all endpoints.
5. Launch Flight proxy server that rewrites endpoints to point to the local port forwards.

This should always be used as a context manager so that all connections are closed when done.

Parameters:

Name	Type	Description	Default
ssh_dest	str	Remote SSH destination (see ssh(1) for form).	required
arrakis_server	str | None	Remote Arrakis info server URL. If not specified it will be determined from the ARRAKIS_SERVER env var on the remote side.	None
bind_address	str	Flight proxy server HOST:PORT. Defaults to "localhost:0" (random open port chosen on localhost).	None

Source code in arrakis/proxy.py

@contextlib.contextmanager
def ssh_proxy(
    ssh_dest: str,
    arrakis_server: str | None = None,
    bind_address: str | None = None,
):
    """Create Flight proxy server over SSH.

    This is done by:

    0. Make SSH connection to the remote host.
    1. Determine initial Flight info server URL on the remote side.
    2. Retrieve all known endpoints from the remote info server.
    3. Set up local ssh port forwards to all endpoints.
    4. Launch Flight proxy server that rewrites endpoints to point to
       the local port forwards.

    This should always be used as a context manager so that all
    connections are closed when done.

    Parameters
    ----------
    ssh_dest : str
        Remote SSH destination (see ssh(1) for form).
    arrakis_server : str | None
        Remote Arrakis info server URL. If not specified it will be
        determined from the ARRAKIS_SERVER env var on the remote side.
    bind_address : str
        Flight proxy server HOST:PORT. Defaults to "localhost:0"
        (random open port chosen on localhost).

    """
    logger.info("creating ssh proxy via %s", ssh_dest)

    # check/hold the requested server port
    bind_socket = socket.socket()
    if bind_address is None:
        bind_address = "localhost:0"
    bind_host, bind_port = bind_address.split(":")
    try:
        bind_socket.bind((bind_host, int(bind_port)))
    except OSError:
        msg = f"local address already in use: {bind_address}"
        raise OSError(msg) from None
    _, bind_port = bind_socket.getsockname()

    # initiate the ssh connection to the host as a context manager, so
    # that the connection is properly shut down if there are any
    # errors during setup.
    with SSHConnection(ssh_dest) as ssh:
        # if not specified, try to determine the remote server location
        # from the ARRAKIS_SERVER env var on the remote host
        if arrakis_server is None:
            logger.debug("resolving remote ARRAKIS_SERVER...")
            arrakis_server = (
                ssh.exec(
                    shell_cmd=["printenv", "ARRAKIS_SERVER"],
                    capture_output=True,
                    check=True,
                )
                .stdout.decode()
                .strip()
            )
            if not arrakis_server:
                msg = "Could not determine remote ARRAKIS_SERVER."
                raise ValueError(msg)

        # create an initial forward to the info server so that we can
        # query for the endpoint information
        backend_hostport = ssh.forward_port(parse_arrakis_url(arrakis_server).netloc)

        logger.debug("retrieving endpoints...")
        backend_url = f"grpc://{backend_hostport}"
        endpoints = Client(backend_url).endpoints()

        # setup forwards for all endpoints
        logger.debug("creating endpoint port forwarding...")
        for endpoint in endpoints:
            remote_hostport = urlparse(endpoint).netloc
            ssh.forward_port(remote_hostport)

        # start the proxy server
        logger.debug("starting Flight proxy server...")
        proxy_url = f"grpc://{bind_host}:{bind_port}"
        bind_socket.close()
        server = FlightProxy(proxy_url, backend_url, ssh.forward_map)

        def _run_until_shutdown(server):
            server.serve()

        executor = ThreadPoolExecutor(max_workers=1)
        executor.submit(_run_until_shutdown, server)

        # yield context manager
        try:
            yield proxy_url
        finally:
            logger.info("shutting down ssh proxy...")
            server.shutdown()
            executor.shutdown()
            logger.info("done.")