From a91571dd222991a13ceeac59a8f7a4c75e7a92e4 Mon Sep 17 00:00:00 2001
From: Connor Dibble <cddibble@gmail.com>
Date: Wed, 3 Nov 2021 22:47:57 +0100
Subject: [PATCH] Support Arm64 CPU architectures

---
 tests/test_conda.py | 10 +++++++---
 tljh/installer.py   | 11 +++++++++--
 tljh/traefik.py     | 19 ++++++++++++++-----
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/tests/test_conda.py b/tests/test_conda.py
index a4db013..7632673 100644
--- a/tests/test_conda.py
+++ b/tests/test_conda.py
@@ -11,11 +11,15 @@ import tempfile
 @pytest.fixture(scope='module')
 def prefix():
     """
-    Provide a temporary directory with a conda environment
+    Provide a temporary directory with a mambaforge conda environment
     """
+    # see https://github.com/conda-forge/miniforge/releases
     mambaforge_version = '4.10.3-7'
-    installer_sha256 = "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474"
-    installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-x86_64.sh".format(v=mambaforge_version)
+    if os.uname().machine == 'aarch64':
+        installer_sha256 = "ac95f137b287b3408e4f67f07a284357b1119ee157373b788b34e770ef2392b2"
+    elif os.uname().machine == 'x86_64':
+        installer_sha256 = "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474"
+    installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-{arch}.sh".format(v=mambaforge_version, arch=os.uname().machine)
     with tempfile.TemporaryDirectory() as tmpdir:
         with conda.download_miniconda_installer(installer_url, installer_sha256) as installer_path:
             conda.install_miniconda(installer_path, tmpdir)
diff --git a/tljh/installer.py b/tljh/installer.py
index 3bb4e84..a5177ef 100644
--- a/tljh/installer.py
+++ b/tljh/installer.py
@@ -170,7 +170,14 @@ def ensure_user_environment(user_requirements_txt_file):
     # Install mambaforge using an installer from
     # https://github.com/conda-forge/miniforge/releases
     mambaforge_new_version = '4.10.3-7'
-    installer_sha256 = "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474"
+    # Check system architecture, set appropriate installer checksum
+    if os.uname().machine == 'aarch64':
+        installer_sha256 = "ac95f137b287b3408e4f67f07a284357b1119ee157373b788b34e770ef2392b2"
+    elif os.uname().machine == 'x86_64':
+        installer_sha256 = "fc872522ec427fcab10167a93e802efaf251024b58cc27b084b915a9a73c4474"
+    # Check OS, set appropriate string for conda installer path
+    if os.uname().sysname != 'Linux':
+        raise OSError("TLJH is only supported on Linux platforms.")
     # Then run `mamba --version` to get the conda and mamba versions
     # Keep these in sync with tests/test_conda.py::prefix
     mambaforge_conda_new_version = '4.10.3'
@@ -185,7 +192,7 @@ def ensure_user_environment(user_requirements_txt_file):
     # If no prior miniconda installation is found, we can install a newer version
     else:
         logger.info('Downloading & setting up user environment...')
-        installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-x86_64.sh".format(v=mambaforge_new_version)
+        installer_url = "https://github.com/conda-forge/miniforge/releases/download/{v}/Mambaforge-{v}-Linux-{arch}.sh".format(v=mambaforge_new_version, arch=os.uname().machine)
         with conda.download_miniconda_installer(installer_url, installer_sha256) as installer_path:
             conda.install_miniconda(installer_path, USER_ENV_PREFIX)
         conda_version = '4.10.3'
diff --git a/tljh/traefik.py b/tljh/traefik.py
index 4c3151f..aa4799f 100644
--- a/tljh/traefik.py
+++ b/tljh/traefik.py
@@ -12,13 +12,21 @@ import toml
 from .config import CONFIG_DIR
 from tljh.configurer import load_config, _merge_dictionaries
 
-# FIXME: support more than one platform here
-plat = "linux-amd64"
-traefik_version = "1.7.18"
+# traefik 2.7.x is not supported yet, use v1.7.x for now
+# see: https://github.com/jupyterhub/traefik-proxy/issues/97
+machine = os.uname().machine
+if machine == 'aarch64':
+    plat = "linux-arm64"
+elif machine == 'x86_64':
+    plat = "linux-amd64"
+else:
+    raise OSError(f"Error. Platform: {os.uname().sysname} / {machine} Not supported.")
+traefik_version = "1.7.33"
 
 # record sha256 hashes for supported platforms here
 checksums = {
-    "linux-amd64": "3c2d153d80890b6fc8875af9f8ced32c4d684e1eb5a46d9815337cb343dfd92e"
+    "linux-amd64": "314ffeaa4cd8ed6ab7b779e9b6773987819f79b23c28d7ab60ace4d3683c5935",
+    "linux-arm64": "0640fa665125efa6b598fc08c100178e24de66c5c6035ce5d75668d3dc3706e1"
 }
 
 def checksum_file(path):
@@ -40,7 +48,7 @@ def fatal_error(e):
     giveup=fatal_error
 )
 def ensure_traefik_binary(prefix):
-    """Download and install the traefik binary"""
+    """Download and install the traefik binary to a location identified by a prefix path such as '/opt/tljh/hub/'"""
     traefik_bin = os.path.join(prefix, "bin", "traefik")
     if os.path.exists(traefik_bin):
         checksum = checksum_file(traefik_bin)
@@ -57,6 +65,7 @@ def ensure_traefik_binary(prefix):
         "https://github.com/containous/traefik/releases"
         f"/download/v{traefik_version}/traefik_{plat}"
     )
+
     print(f"Downloading traefik {traefik_version}...")
     # download the file
     response = requests.get(traefik_url)