Rewrite the micromamba Feature (#3)

Co-authored-by: eitsupi <50911393+eitsupi@users.noreply.github.com>

.shellcheckrc

@@ -0,0 +1 @@
+external-sources=true

src/micromamba/NOTES.md

@@ -0,0 +1,20 @@
+## Version number specification
+
+Soft version matching is *not* supported, meaning that `"1"` and `"1.0"` are not
+valid values of the `version` parameter. The full version number must be specified
+like `"1.0.0"`.
+
+## Channels
+
+By default, `micromamba` configures no channels. If you would like to set `conda-forge`
+as a default channel, then use
+
+```json
+"features": {
+  "ghcr.io/eitsupi/mamba-devcontainer-features/micromamba:0.1": {
+    "channels": "conda-forge"
+  }
+}
+```
+
+More generally, `channels` can be a comma-separated list such as "conda-forge,defaults".

src/micromamba/devcontainer-feature.json

@@ -8,13 +8,29 @@
         "version": {
             "type": "string",
             "proposals": [
-                "latest",
-                "1"
+                "latest"
             ],
             "default": "latest",
-            "description": "Select version of micromamba."
+            "description": "Exact version of Micromamba to install, if not latest (must be X.Y.Z)"
+        },
+        "allowReinstall": {
+            "type": "boolean",
+            "default": false,
+            "description": "Reinstall in case Micromamba already exists"
+        },
+        "channels": {
+            "type": "string",
+            "default": "",
+            "proposals": [
+                "conda-forge"
+            ],
+            "description": "Comma separated list of Conda channels to add"
         }
     },
+    "containerEnv": {
+        "MAMBA_ROOT_PREFIX": "/opt/conda",
+        "PATH": "/opt/conda/bin:${PATH}"
+    },
     "installsAfter": [
         "ghcr.io/devcontainers/features/common-utils"
     ]

src/micromamba/install.sh

@@ -1,101 +1,45 @@
 #!/usr/bin/env bash
 
-VERSION=${VERSION:-"latest"}
-
-USERNAME=${USERNAME:-"automatic"}
-
 set -e
 
-# Clean up
-rm -rf /var/lib/apt/lists/*
+# Move to the same directory as this script
+FEATURE_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+cd "${FEATURE_DIR}"
 
-if [ "$(id -u)" -ne 0 ]; then
-    echo -e 'Script must be run as root. Use sudo, su, or add "USER root" to your Dockerfile before running this script.'
-    exit 1
-fi
+# Options
+VERSION=${VERSION:-"latest"}
+ALLOW_REINSTALL=${ALLOWREINSTALL:-"false"}
+IFS=',' read -r -a CHANNELS <<< "$CHANNELS"  # Convert comma-separated list to array
 
-architecture="$(dpkg --print-architecture)"
-if [ "${architecture}" != "amd64" ] && [ "${architecture}" != "arm64" ]; then
-    echo "(!) Architecture $architecture unsupported"
-    exit 1
-fi
+# Constants
+MAMBA_ROOT_PREFIX="/opt/conda"
+micromamba_destination="/usr/local/bin"
 
-# Determine the appropriate non-root user
-if [ "${USERNAME}" = "auto" ] || [ "${USERNAME}" = "automatic" ]; then
-    USERNAME=""
-    POSSIBLE_USERS=("vscode" "node" "codespace" "$(awk -v val=1000 -F ":" '$3==val{print $1}' /etc/passwd)")
-    for CURRENT_USER in "${POSSIBLE_USERS[@]}"; do
-        if id -u "${CURRENT_USER}" >/dev/null 2>&1; then
-            USERNAME=${CURRENT_USER}
-            break
-        fi
-    done
-    if [ "${USERNAME}" = "" ]; then
-        USERNAME=root
-    fi
-elif [ "${USERNAME}" = "none" ] || ! id -u ${USERNAME} >/dev/null 2>&1; then
-    USERNAME=root
-fi
+# shellcheck source=./utils.sh
+source ./utils.sh
 
-apt_get_update() {
-    if [ "$(find /var/lib/apt/lists/* | wc -l)" = "0" ]; then
-        echo "Running apt-get update..."
-        apt-get update -y
-    fi
-}
+# Note: The apt-cache is cleared on-demand.
+# Thus we don't need here "rm -rf /var/lib/apt/lists/*".
 
-# Checks if packages are installed and installs them if not
-check_packages() {
-    if ! dpkg -s "$@" >/dev/null 2>&1; then
-        apt_get_update
-        apt-get -y install --no-install-recommends "$@"
-    fi
-}
+USERNAME="${USERNAME:-"${_REMOTE_USER:-"automatic"}"}"
+detect_user USERNAME
 
-check_git() {
-    if [ ! -x "$(command -v git)" ]; then
-        check_packages git
-    fi
+require_running_as_root
+
+ensure_download_prerequisites() {
+    # This is the only place we need to use apt, so we can scope clean_up_apt tightly:
+    check_packages curl ca-certificates bzip2
 }
 
-find_version_from_git_tags() {
-    local variable_name=$1
-    local requested_version=${!variable_name}
-    if [ "${requested_version}" = "none" ]; then return; fi
-    local repository=$2
-    local prefix=${3:-"tags/v"}
-    local separator=${4:-"."}
-    local last_part_optional=${5:-"false"}
-    if [ "$(echo "${requested_version}" | grep -o "." | wc -l)" != "2" ]; then
-        local escaped_separator=${separator//./\\.}
-        local last_part
-        if [ "${last_part_optional}" = "true" ]; then
-            last_part="(${escaped_separator}[0-9]+)*?"
-        else
-            last_part="${escaped_separator}[0-9]+"
-        fi
-        local regex="${prefix}\\K[0-9]+${escaped_separator}[0-9]+${last_part}$"
-        local version_list
-        check_git
-        check_packages ca-certificates
-        version_list="$(git ls-remote --tags "${repository}" | grep -oP "${regex}" | tr -d ' ' | tr "${separator}" "." | sort -rV)"
-        if [ "${requested_version}" = "latest" ] || [ "${requested_version}" = "current" ] || [ "${requested_version}" = "lts" ]; then
-            declare -g "${variable_name}"="$(echo "${version_list}" | head -n 1)"
-        else
-            set +e
-            declare -g "${variable_name}"="$(echo "${version_list}" | grep -E -m 1 "^${requested_version//./\\.}([\\.\\s]|$)")"
-            set -e
-        fi
-    fi
-    if [ -z "${!variable_name}" ] || ! echo "${version_list}" | grep "^${!variable_name//./\\.}$" >/dev/null 2>&1; then
-        echo -e "Invalid ${variable_name} value: ${requested_version}\nValid values:\n${version_list}" >&2
-        exit 1
-    fi
-    echo "${variable_name}=${!variable_name}"
+download_with_curl() {
+    local url=$1
+    local destination=$2
+    curl -sL "${url}" | tar -xj -C "${destination}" --strip-components=1 bin/micromamba
 }
 
 install_micromamba() {
     local version=$1
+    local destination=$2
     local arch
     local url
     arch="$(uname -m)"
@@ -104,19 +48,116 @@ install_micromamba() {
     fi
     url="https://micro.mamba.pm/api/micromamba/linux-${arch}/${version}"
 
-    check_packages curl ca-certificates bzip2
-    echo "Downloading micromamba..."
-    curl -sL "${url}" | tar -xj -C /usr/local/bin/ --strip-components=1 bin/micromamba
+    echo "Installing prerequisites for downloading micromamba..."
+    ensure_download_prerequisites
+    echo "Downloading micromamba from ${url}..."
+    download_with_curl "${url}" "${destination}"
+    echo "Micromamba download complete."
+}
+
+run_as_user() {
+    su "${USERNAME}" "${@}"
+}
+
+micromamba_as_user() {
+    run_as_user bash -c "micromamba $*"
+}    
+
+add_conda_group() {
+    if ! cat /etc/group | grep -e "^conda:" > /dev/null 2>&1; then
+        groupadd -r conda
+    fi
+}
+
+initialize_root_prefix() {
+    mkdir -p "${MAMBA_ROOT_PREFIX}/conda-meta"
+    touch "${MAMBA_ROOT_PREFIX}/conda-meta/history"
+    add_conda_group
+    usermod -a -G conda "${USERNAME}"
+    chown -R "${USERNAME}:conda" "${MAMBA_ROOT_PREFIX}"
+    chmod -R g+r+w "${MAMBA_ROOT_PREFIX}"
+    find "${MAMBA_ROOT_PREFIX}" -type d -print0 | xargs -n 1 -0 chmod g+s
+}
+
+add_channels() {
+    # Channels source: <https://docs.anaconda.com/anaconda/user-guide/tasks/using-repositories/>
+    anaconda_channels=("defaults" "main" "r" "msys2" "free" "mro" "mro-archive" "archive" "pro" "anaconda-extras" "anaconda")
+    for channel in "${CHANNELS[@]}"; do
+        if [ ! -z "${channel}" ]; then
+            echo "Adding channel ${channel}"
+            micromamba_as_user config append channels "${channel}"
+            if [[ " ${anaconda_channels[*]} " =~ " ${channel} " ]]; then
+                make_anaconda_repository_warning
+            fi
+        fi
+    done
+}
+
+make_anaconda_repository_warning() {
+    # Original source:
+    # <https://github.com/devcontainers/features/blob/baf47e22b0c3dc5b418ac57aae2e750d14bbc9a3/src/conda/install.sh#L104-L119>
+
+    # Display a notice on conda when not running in GitHub Codespaces
+    mkdir -p /usr/local/etc/vscode-dev-containers
+    cat << 'EOF' > /usr/local/etc/vscode-dev-containers/conda-notice.txt
+When using "micromamba" from outside of GitHub Codespaces, note the Anaconda repository contains
+restrictions on commercial use that may impact certain organizations. See https://aka.ms/ghcs-conda
+EOF
+
+    notice_script="$(cat << 'EOF'
+if [ -t 1 ] && [ "${IGNORE_NOTICE}" != "true" ] && [ "${TERM_PROGRAM}" = "vscode" ] && [ "${CODESPACES}" != "true" ] && [ ! -f "$HOME/.config/vscode-dev-containers/conda-notice-already-displayed" ]; then
+    cat "/usr/local/etc/vscode-dev-containers/conda-notice.txt"
+    mkdir -p "$HOME/.config/vscode-dev-containers"
+    ((sleep 10s; touch "$HOME/.config/vscode-dev-containers/conda-notice-already-displayed") &)
+fi
+EOF
+)"
+
+    if [ -f "/etc/zsh/zshrc" ]; then
+        echo "${notice_script}" | tee -a /etc/zsh/zshrc
+    fi
+
+    if [ -f "/etc/bash.bashrc" ]; then
+        echo "${notice_script}" | tee -a /etc/bash.bashrc
+    fi
 }
 
 export DEBIAN_FRONTEND=noninteractive
 
-# Soft version matching
-find_version_from_git_tags VERSION "https://github.com/mamba-org/mamba" "tags/micromamba-"
+ensure_path_for_login_shells
+
+if [ "${ALLOW_REINSTALL}" = "false" ]; then
+    if type micromamba > /dev/null 2>&1; then
+        echo "Detected existing micromamba: $(micromamba --version)."
+        echo "The allowReinstall argument is false, so not overwriting."
+        skip_install="true"
+    fi
+fi
+
+if [ "${skip_install}" != "true" ]; then
+    install_micromamba "${VERSION}" "${micromamba_destination}"
+    echo "Micromamba executable installed."
+fi
+
+initialize_root_prefix
+
+add_channels
+
+echo "Setting channel_priority to strict"
+micromamba_as_user config set channel_priority strict
+
+echo "Initializing Bash shell"
+micromamba_as_user shell init --shell=bash
+su -c "if ! grep -q 'micromamba activate # added by micromamba devcontainer feature' ~/.bashrc; then echo 'micromamba activate # added by micromamba devcontainer feature' >> ~/.bashrc; fi" - "${USERNAME}"
+
+if type zsh > /dev/null 2>&1; then
+    echo "Initializing zsh shell"
+    micromamba_as_user shell init --shell=zsh
+    su -c "if ! grep -q 'micromamba activate # added by micromamba devcontainer feature' ~/.zshrc; then echo 'micromamba activate # added by micromamba devcontainer feature' >> ~/.zshrc; fi" - "${USERNAME}"
+fi
 
-install_micromamba "${VERSION##*-}"
+echo "Micromamba configured."
 
-# Clean up
-rm -rf /var/lib/apt/lists/*
+clean_up_apt_if_updated
 
 echo "Done!"

src/micromamba/utils.sh

@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+
+# In case we need to use apt-get we cannot trust the cache since it might be old.
+# To avoid any unnecessary operations, we keep track of the cache state here.
+apt_cache_state="unaccessed"
+# Other possible values of apt_cache_state are "updated" and "clean".
+# If apt-get is never run, then this remains as "unaccessed".
+# The only functions which directly modify this state are:
+#   - clean_up_apt
+#   - apt_get_update
+# It is indirectly modified by:
+#   - clean_up_apt_if_updated
+#   - check_packages
+# All apt related invocations should use only these functions.
+# If apt-get is required, then the progression is:
+#   - "unaccessed"
+#   # clean_up_apt
+#   - "clean"
+#   # apt-get update
+#   - "updated"
+#   # apt-get install ...  # one or more times, via check_packages
+#   # clean_up_apt  # via the clean_up_apt_if_updated function
+#   - "clean"
+# and then the script exits.
+
+clean_up_apt() {
+    rm -rf /var/lib/apt/lists/*
+    apt_cache_state="clean"
+}
+
+clean_up_apt_if_updated() {
+    if [ "${apt_cache_state}" = "updated" ]; then
+        clean_up_apt
+    fi
+}
+
+apt_get_update() {
+    if [ "${apt_cache_state}" = "unaccessed" ]; then
+        clean_up_apt
+    fi
+    if [ "${apt_cache_state}" = "clean" ]; then
+        echo "Running apt-get update..."
+        apt-get update -y
+        apt_cache_state="updated"
+    fi
+}
+
+# Checks if packages are installed and installs them if not
+check_packages() {
+    if ! dpkg -s "$@" >/dev/null 2>&1; then
+        apt_get_update
+        apt-get -y install --no-install-recommends "$@"
+    fi
+}
+
+require_running_as_root() {
+    local error_message="${1:-Script must be run as root. Use sudo, su, or add \"USER root\" to your Dockerfile before running this script.}"
+    if [ "$(id -u)" -ne 0 ]; then
+        echo -e "${error_message}"
+        exit 1
+    fi
+}
+
+# Source:
+# <https://github.com/devcontainers/features/blob/7b009e661f13085629b19fc157b577916587f6bc/src/nix/utils.sh#L67-L83>
+# If in automatic mode, determine if a user already exists, if not use root
+detect_user() {
+    local user_variable_name=${1:-username}
+    local possible_users=("vscode" "node" "codespace" "$(awk -v val=1000 -F ":" '$3==val{print $1}' /etc/passwd)")
+    if [ "${!user_variable_name}" = "auto" ] || [ "${!user_variable_name}" = "automatic" ]; then
+        declare -g ${user_variable_name}=""
+        for current_user in ${possible_users[@]}; do
+            if id -u "${current_user}" > /dev/null 2>&1; then
+                declare -g ${user_variable_name}="${current_user}"
+                break
+            fi
+        done
+    fi
+    if [ "${!user_variable_name}" = "" ] || [ "${!user_variable_name}" = "none" ] || ! id -u "${!user_variable_name}" > /dev/null 2>&1; then
+        declare -g ${user_variable_name}=root
+    fi
+}
+
+ensure_path_for_login_shells() {
+    # Ensure that login shells get the correct path if the user updated the PATH using ENV.
+    rm -f /etc/profile.d/00-restore-env.sh
+    echo "export PATH=${PATH//$(sh -lc 'echo $PATH')/\$PATH}" > /etc/profile.d/00-restore-env.sh
+    chmod +x /etc/profile.d/00-restore-env.sh
+}