From e71b33dbe0f538cbe7f791eb6840172c710a46a5 Mon Sep 17 00:00:00 2001
From: Hugo Shaka <hugo.hervieux@goteleport.com>
Date: Wed, 6 Mar 2024 13:22:28 -0500
Subject: [PATCH] Query target version in discovery script

---
 api/types/installers/installer.sh.tmpl | 87 +++++++++++++++++++++++---
 1 file changed, 80 insertions(+), 7 deletions(-)

diff --git a/api/types/installers/installer.sh.tmpl b/api/types/installers/installer.sh.tmpl
index 9439e5fc9a53a..2984370aa0be6 100644
--- a/api/types/installers/installer.sh.tmpl
+++ b/api/types/installers/installer.sh.tmpl
@@ -3,6 +3,32 @@
 
 set -eu
 
+upgrade_endpoint="{{ .PublicProxyAddr }}/v1/webapi/automaticupgrades/channel/default"
+
+# upgrade_endpoint_fetch loads the specified value from the upgrade endpoint. the only
+# currently supported values are 'version' and 'critical'.
+upgrade_endpoint_fetch() {
+    host_path="${upgrade_endpoint}/${1}"
+
+    if sf_output="$(curl --proto '=https' --tlsv1.2 -sSf "https://${host_path}")"; then
+        # emit output with empty lines and extra whitespace removed
+        echo "$sf_output" | grep -v -e '^[[:space:]]*$' | awk '{$1=$1};1'
+        return 0
+    else
+        return 1
+    fi
+}
+
+# get_target_version loads the current value of the /version endpoint.
+get_target_version() {
+    if tv_output="$(upgrade_endpoint_fetch version)"; then
+        # emit version string with leading 'v' removed if one is present
+        echo "${tv_output#v}"
+        return 0
+    fi
+    return 1
+}
+
 on_ec2() {
   IMDS_TOKEN=$(curl -m5 -sS -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 300")
   [ -z "$IMDS_TOKEN" ] && return 1
@@ -28,11 +54,8 @@ on_gcp() {
   # shellcheck disable=SC1091
   . /etc/os-release
 
-  PACKAGE_LIST="{{ .TeleportPackage }} jq"
-  # shellcheck disable=SC2050
-  if [ "{{ .AutomaticUpgrades }}" = "true" ]; then
-    PACKAGE_LIST="${PACKAGE_LIST} {{ .TeleportPackage }}-updater"
-  fi
+  TELEPORT_PACKAGE="{{ .TeleportPackage }}"
+  TELEPORT_UPDATER_PACKAGE="{{ .TeleportPackage }}"
 
   if [ "$ID" = "debian" ] || [ "$ID" = "ubuntu" ]; then
     # old versions of ubuntu require that keys get added by `apt-key add`, without
@@ -48,7 +71,24 @@ on_gcp() {
       echo "deb [signed-by=/usr/share/keyrings/teleport-archive-keyring.asc]  https://apt.releases.teleport.dev/${ID?} ${VERSION_CODENAME?} {{ .RepoChannel }}" | sudo tee /etc/apt/sources.list.d/teleport.list >/dev/null
     fi
     sudo apt-get update
-    sudo apt-get install -y ${PACKAGE_LIST}
+
+    if [ "{{ .AutomaticUpgrades }}" = "true" ]; then
+      # automatic upgrades
+      if ! target_version="$(get_target_version)"; then
+        # error getting the target version
+        sudo apt-get install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      elif [ "$target_version" == "none" ]; then
+        # no target version advertised
+        sudo apt-get install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      else
+        # successfully retrieved target version
+        sudo apt-get install -y "$TELEPORT_PACKAGE=$target_version" jq "$TELEPORT_UPDATER_PACKAGE=$target_version"
+      fi
+    else
+      # no automatic upgrades
+      sudo apt-get install -y "$TELEPORT_PACKAGE"
+    fi
+
   elif [ "$ID" = "amzn" ] || [ "$ID" = "rhel" ]; then
     if [ "$ID" = "rhel" ]; then
       VERSION_ID=${VERSION_ID//\.*/} # convert version numbers like '7.2' to only include the major version
@@ -56,7 +96,24 @@ on_gcp() {
     sudo yum install -y yum-utils
     sudo yum-config-manager --add-repo \
       "$(rpm --eval "https://yum.releases.teleport.dev/$ID/$VERSION_ID/Teleport/%{_arch}/{{ .RepoChannel }}/teleport.repo")"
-    sudo yum install -y ${PACKAGE_LIST}
+
+    if [ "{{ .AutomaticUpgrades }}" = "true" ]; then
+      # automatic upgrades
+      if ! target_version="$(get_target_version)"; then
+        # error getting the target version
+        sudo yum install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      elif [ "$target_version" == "none" ]; then
+        # no target version advertised
+        sudo yum install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      else
+        # successfully retrieved target version
+        sudo yum install -y "$TELEPORT_PACKAGE-$target_version" jq "$TELEPORT_UPDATER_PACKAGE-$target_version"
+      fi
+    else
+      # no automatic upgrades
+      sudo yum install -y "$TELEPORT_PACKAGE"
+    fi
+
   elif [ "$ID" = "sles" ] || [ "$ID" = "opensuse-tumbleweed" ] || [ "$ID" = "opensuse-leap" ]; then
     if [ "$ID" = "opensuse-tumbleweed" ]; then
       VERSION_ID="15" # tumbleweed uses dated VERSION_IDs like 20230702
@@ -67,6 +124,22 @@ on_gcp() {
     sudo zypper --non-interactive addrepo "$(rpm --eval "https://zypper.releases.teleport.dev/sles/$VERSION_ID/Teleport/%{_arch}/{{ .RepoChannel }}/teleport.repo")"
     sudo zypper --gpg-auto-import-keys refresh
     sudo zypper --non-interactive install ${PACKAGE_LIST}
+    if [ "{{ .AutomaticUpgrades }}" = "true" ]; then
+      # automatic upgrades
+      if ! target_version="$(get_target_version)"; then
+        # error getting the target version
+        sudo zypper --non-interactive install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      elif [ "$target_version" == "none" ]; then
+        # no target version advertised
+        sudo zypper --non-interactive install -y "$TELEPORT_PACKAGE" jq "$TELEPORT_UPDATER_PACKAGE"
+      else
+        # successfully retrieved target version
+        sudo zypper --non-interactive install -y "$TELEPORT_PACKAGE-$target_version" jq "$TELEPORT_UPDATER_PACKAGE-$target_version"
+      fi
+    else
+      # no automatic upgrades
+      sudo zypper --non-interactive install -y "$TELEPORT_PACKAGE"
+    fi
   else
     echo "Unsupported distro: $ID"
     exit 1