Merge pull request #724 from splitgraph/tunnel-issues-2-CU-38f3r2e

Download rathole client binary for platform if necessary

Author: neumark

splitgraph/cloud/tunnel_client.py

@@ -1,8 +1,13 @@
+import hashlib
 import os
+import platform
 import subprocess
+from io import BytesIO
 from os import path
 from typing import Optional
+from zipfile import ZipFile
 
+import requests
 from click import echo
 
 RATHOLE_CLIENT_CONFIG_FILENAME = "rathole-client.toml"
@@ -24,6 +29,45 @@
 
 """
 
+# Download URL and SHA256 hash of rathole build ZIP archive for each supported platform.
+RATHOLE_BUILDS = {
+    "Darwin": (
+        "https://github.com/rapiz1/rathole/releases/download/v0.4.4/rathole-x86_64-apple-darwin.zip",
+        "c1e6d0a41a0af8589303ab6940937d9183b344a62283ff6033a17e82c357ce17",
+    ),
+    "Windows": (
+        "https://github.com/rapiz1/rathole/releases/download/v0.4.4/rathole-x86_64-pc-windows-msvc.zip",
+        "92cc3feb57149c0b4dba7ec198dbda26c4831cde0a7c74a7d9f51e0002f65ead",
+    ),
+    "Linux-x86_64-glibc": (
+        "https://github.com/rapiz1/rathole/releases/download/v0.4.4/rathole-x86_64-unknown-linux-gnu.zip",
+        "fef39ed9d25e944711e2a27d5a9c812163ab184bf3f703827fca6bbf54504fbf",
+    ),
+    "Linux-x86_64-musl": (
+        "https://github.com/rapiz1/rathole/releases/download/v0.4.4/rathole-x86_64-unknown-linux-musl.zip",
+        "fc6b0a57727383a1491591f8e9ee76b1e0e25ecf7c2736b803d8f4411f651a15",
+    ),
+}
+
+
+def get_sha256(stream):
+    return hashlib.sha256(stream.read()).hexdigest()
+
+
+def get_rathole_build_key():
+    system = platform.system()
+    # Currently only x86_64 macos builds exist for Windows and MacOS.
+    # It works on Apple Silicon due to Rosetta 2.
+    if system in ["Windows", "Darwin"]:
+        return system
+    if system == "Linux":
+        # python 3.8 sometimes reports '' instead of 'musl' for musl libc (https://bugs.python.org/issue43248)
+        return "Linux-%s-%s" % (
+            platform.machine(),
+            "glibc" if platform.libc_ver()[0] == "glibc" else "musl",
+        )
+    return None
+
 
 def get_rathole_client_config(
     tunnel_connect_address: str,
@@ -51,8 +95,31 @@ def get_config_dir():
     return os.path.dirname(get_singleton(CONFIG, "SG_CONFIG_FILE"))
 
 
-def get_rathole_client_binary_path():
-    return os.path.join(get_config_dir(), "rathole")
+def get_rathole_client_binary_path(config_dir: Optional[str] = None) -> str:
+    filename = "rathole.exe" if get_rathole_build_key() == "Windows" else "rathole"
+    return os.path.join(config_dir or get_config_dir(), filename)
+
+
+def get_config_filename(config_dir: Optional[str] = None) -> str:
+    return path.join(config_dir or get_config_dir(), RATHOLE_CLIENT_CONFIG_FILENAME)
+
+
+def download_rathole_binary(
+    build: Optional[str] = None, rathole_path: Optional[str] = None
+) -> None:
+    rathole_binary_path = rathole_path or get_rathole_client_binary_path()
+    (url, sha256) = RATHOLE_BUILDS.get(build or get_rathole_build_key(), ("", ""))
+    if not url:
+        raise Exception("No rathole build found for this architecture")
+    content = BytesIO(requests.get(url).content)
+    assert get_sha256(content) == sha256
+    content.seek(0)
+    zipfile = ZipFile(content)
+    assert len(zipfile.filelist) == 1
+    assert zipfile.filelist[0].filename == path.basename(rathole_binary_path)
+    zipfile.extract(path.basename(rathole_binary_path), path.dirname(rathole_binary_path))
+    if get_rathole_build_key() != "Windows":
+        os.chmod(rathole_binary_path, 0o500)
 
 
 def write_rathole_client_config(
@@ -64,7 +131,7 @@ def write_rathole_client_config(
     tls_hostname: Optional[str],
 ) -> str:
     # If a specific root CA file is used (eg: for self-signed hosts), reference
-    # it in the rathole client config.
+    # it in the rathole client config. Otherwise use system default trust root.
     trusted_root = os.environ.get("REQUESTS_CA_BUNDLE") or os.environ.get("SSL_CERT_FILE")
     rathole_client_config = get_rathole_client_config(
         tunnel_connect_address=f"{tunnel_connect_host}:{tunnel_connect_port}",
@@ -74,13 +141,23 @@ def write_rathole_client_config(
         section_id=section_id,
         trusted_root=trusted_root,
     )
-    config_filename = path.join(get_config_dir(), RATHOLE_CLIENT_CONFIG_FILENAME)
+    config_filename = get_config_filename()
     with open(config_filename, "w") as f:
         f.write(rathole_client_config)
     return config_filename
 
 
-def launch_rathole_client(rathole_client_config_path: str) -> None:
+def launch_rathole_client(
+    rathole_client_binary_path: Optional[str] = None,
+    rathole_client_config_path: Optional[str] = None,
+) -> None:
+    binary_path = rathole_client_binary_path or get_rathole_client_binary_path()
+    if not path.isfile(binary_path):
+        download_rathole_binary()
     echo("launching rathole client")
-    command = [get_rathole_client_binary_path(), "--client", rathole_client_config_path]
+    command = [
+        binary_path,
+        "--client",
+        rathole_client_config_path or get_config_filename(),
+    ]
     subprocess.check_call(command)

splitgraph/commandline/cloud.py

@@ -1238,7 +1238,7 @@ def start_repository_tunnel(
     section_id = f"{repository.namespace}/{repository.repository}"
     local_address = f"{external.params['host']}:{external.params['port']}"
 
-    rathole_client_config_path = write_rathole_client_config(
+    write_rathole_client_config(
         section_id,
         secret_token,
         tunnel_connect_host,
@@ -1247,7 +1247,7 @@ def start_repository_tunnel(
         tunnel_connect_host,
     )
 
-    launch_rathole_client(rathole_client_config_path)
+    launch_rathole_client()
 
 
 def start_ephemeral_tunnel(remote: str, local_address: str) -> None:
@@ -1262,7 +1262,7 @@ def start_ephemeral_tunnel(remote: str, local_address: str) -> None:
         private_address_port,
     ) = client.provision_ephemeral_tunnel()
 
-    rathole_client_config_path = write_rathole_client_config(
+    write_rathole_client_config(
         private_address_host,
         secret_token,
         tunnel_connect_host,
@@ -1273,7 +1273,7 @@ def start_ephemeral_tunnel(remote: str, local_address: str) -> None:
     click.echo(
         f"To connect to {local_address} from Splitgraph, use the following connection parameters:\nHost: {private_address_host}\nPort: {private_address_port}"
     )
-    launch_rathole_client(rathole_client_config_path)
+    launch_rathole_client()
 
 
 @click.command("tunnel")

splitgraph/config/management.py

@@ -1,6 +1,7 @@
 import logging
 import os
 from pathlib import Path
+from typing import cast
 
 from splitgraph.config.config import patch_config
 from splitgraph.config.export import overwrite_config
@@ -12,7 +13,11 @@ def patch_and_save_config(config, patch):
     config_path = config["SG_CONFIG_FILE"]
     if not config_path:
         # Default to creating a config in the user's homedir rather than local.
-        config_dir = Path(os.environ["HOME"]) / Path(HOME_SUB_DIR)
+        homedir = os.environ.get("HOME")
+        # on Windows, HOME is not a standard env var
+        if homedir is None and os.name == "nt":
+            homedir = f"{os.environ['HOMEDRIVE']}{os.environ['HOMEPATH']}"
+        config_dir = Path(cast(str, homedir)) / Path(HOME_SUB_DIR)
         config_path = config_dir / Path(".sgconfig")
         logging.debug("No config file detected, creating one at %s" % config_path)
         config_dir.mkdir(exist_ok=True, parents=True)

test/splitgraph/commandline/test_cloud.py

@@ -27,6 +27,8 @@
 from httpretty.core import HTTPrettyRequest
 
 from splitgraph.__version__ import __version__
+from splitgraph.cloud.project.models import External
+from splitgraph.cloud.tunnel_client import get_config_filename
 from splitgraph.commandline import cli
 from splitgraph.commandline.cloud import (
     add_c,
@@ -37,6 +39,7 @@
     register_c,
     sql_c,
     stub_c,
+    tunnel_c,
 )
 from splitgraph.config import create_config_dict
 from splitgraph.config.config import patch_config
@@ -624,3 +627,45 @@ def test_commandline_stub(snapshot):
         assert result.exit_code == 0
 
         snapshot.assert_match(result.stdout, "sgr_cloud_stub.yml")
+
+
+def test_rathole_client_config():
+    runner = CliRunner(mix_stderr=False)
+    external = External(
+        tunnel=True, plugin="asdf", params={"host": "127.0.0.1", "port": 5432}, tables={}
+    )
+
+    def mock_provision_tunnel(a, b):
+        print("asdf", a, b)
+        return ("foo", "bar", 1)
+
+    with patch(
+        "splitgraph.commandline.cloud._get_external_from_yaml", return_value=(external,)
+    ), patch(
+        "splitgraph.cloud.GQLAPIClient.provision_repository_tunnel",
+        new_callable=PropertyMock,
+        return_value=mock_provision_tunnel,
+    ), patch(
+        "splitgraph.commandline.cloud.launch_rathole_client", return_value=None
+    ):
+        result = runner.invoke(
+            tunnel_c,
+            [
+                "test/repo",
+            ],
+            catch_exceptions=False,
+        )
+        assert result.exit_code == 0
+        with open(get_config_filename(), "r") as f:
+            non_empty_lines = [line.strip() for line in f if line.strip() != ""]
+            assert non_empty_lines == [
+                "[client]",
+                'remote_addr = "bar:1"',
+                "[client.transport]",
+                'type = "tls"',
+                "[client.transport.tls]",
+                'hostname = "bar"',
+                '[client.services."test/repo"]',
+                'local_addr = "127.0.0.1:5432"',
+                'token = "foo"',
+            ]