update: refactor and fix logic in check_for_upgrade.sh (#8939)

1 files changed, 72 insertions, 46 deletions

diff --git a/tools/check_for_upgrade.sh b/tools/check_for_upgrade.sh
index c8dedcf77..872f5e820 100644
--- a/tools/check_for_upgrade.sh
+++ b/tools/check_for_upgrade.sh
@@ -1,62 +1,88 @@
-#!/usr/bin/env zsh
+# Migrate .zsh-update file to $ZSH_CACHE_DIR
+if [[ -f ~/.zsh-update && ! -f "${ZSH_CACHE_DIR}/.zsh-update" ]]; then
+    mv ~/.zsh-update "${ZSH_CACHE_DIR}/.zsh-update"
+fi
+
+# Cancel update if:
+# - the automatic update is disabled.
+# - the current user doesn't have write permissions nor owns the $ZSH directory.
+# - git is unavailable on the system.
+if [[ "$DISABLE_AUTO_UPDATE" = true ]] \
+   || [[ ! -w "$ZSH" || ! -O "$ZSH" ]] \
+   || ! command -v git &>/dev/null; then
+    return
+fi
 
-zmodload zsh/datetime
 
-function _current_epoch() {
-  echo $(( $EPOCHSECONDS / 60 / 60 / 24 ))
+function current_epoch() {
+    zmodload zsh/datetime
+    echo $(( EPOCHSECONDS / 60 / 60 / 24 ))
 }
 
-function _update_zsh_update() {
-  echo "LAST_EPOCH=$(_current_epoch)" >! ${ZSH_CACHE_DIR}/.zsh-update
+function update_last-updated_file() {
+    echo "LAST_EPOCH=$(current_epoch)" >! "${ZSH_CACHE_DIR}/.zsh-update"
 }
 
-function _upgrade_zsh() {
-  env ZSH=$ZSH sh $ZSH/tools/upgrade.sh
-  # update the zsh file
-  _update_zsh_update
+function update_ohmyzsh() {
+    ZSH="$ZSH" sh "$ZSH/tools/upgrade.sh"
+    update_last-updated_file
 }
 
-epoch_target=$UPDATE_ZSH_DAYS
-if [[ -z "$epoch_target" ]]; then
-  # Default to old behavior
-  epoch_target=13
-fi
+() {
+    emulate -L zsh
+
+    local epoch_target mtime option LAST_EPOCH
+
+    # Remove lock directory if older than a day
+    zmodload zsh/datetime
+    zmodload -F zsh/stat b:zstat
+    if mtime=$(zstat +mtime "$ZSH/log/update.lock" 2>/dev/null); then
+        if (( (mtime + 3600 * 24) < EPOCHSECONDS )); then
+            command rm -rf "$ZSH/log/update.lock"
+        fi
+    fi
 
-# Cancel upgrade if the current user doesn't have write permissions for the
-# oh-my-zsh directory.
-[[ -w "$ZSH" ]] || return 0
+    # Check for lock directory
+    if ! command mkdir "$ZSH/log/update.lock" 2>/dev/null; then
+        return
+    fi
 
-# Cancel upgrade if git is unavailable on the system
-whence git >/dev/null || return 0
+    # Remove lock directory on exit. `return 1` is important for when trapping a SIGINT:
+    #  The return status from the function is handled specially. If it is zero, the signal is
+    #  assumed to have been handled, and execution continues normally. Otherwise, the shell
+    #  will behave as interrupted except that the return status of the trap is retained.
+    trap "rm -rf '$ZSH/log/update.lock'; return 1" EXIT INT QUIT
 
-if mkdir "$ZSH/log/update.lock" 2>/dev/null; then
-  if [ -f ${ZSH_CACHE_DIR}/.zsh-update ]; then
-    . ${ZSH_CACHE_DIR}/.zsh-update
+    # Create or update .zsh-update file if missing or malformed
+    if ! source "${ZSH_CACHE_DIR}/.zsh-update" 2>/dev/null || [[ -z "$LAST_EPOCH" ]]; then
+        update_last-updated_file
+        return
+    fi
 
-    if [[ -z "$LAST_EPOCH" ]]; then
-      _update_zsh_update
-      rmdir $ZSH/log/update.lock # TODO: fix later
-      return 0
+    # Number of days before trying to update again
+    epoch_target=${UPDATE_ZSH_DAYS:-13}
+    # Test if enough time has passed until the next update
+    if (( ( $(current_epoch) - $LAST_EPOCH ) < $epoch_target )); then
+        return
     fi
 
-    epoch_diff=$(($(_current_epoch) - $LAST_EPOCH))
-    if [ $epoch_diff -gt $epoch_target ]; then
-      if [ "$DISABLE_UPDATE_PROMPT" = "true" ]; then
-        _upgrade_zsh
-      else
-        echo "[Oh My Zsh] Would you like to update? [Y/n]: \c"
-        read line
-        if [[ "$line" == Y* ]] || [[ "$line" == y* ]] || [ -z "$line" ]; then
-          _upgrade_zsh
-        else
-          _update_zsh_update
-        fi
-      fi
+    # Ask for confirmation before updating unless disabled
+    if [[ "$DISABLE_UPDATE_PROMPT" = true ]]; then
+        update_ohmyzsh
+    else
+        # input sink to swallow all characters typed before the prompt
+        # and add a newline if there wasn't one after characters typed
+        while read -t -k 1 option; do true; done
+        [[ "$option" != ($'\n'|"") ]] && echo
+
+        echo -n "[oh-my-zsh] Would you like to update? [Y/n] "
+        read -r -k 1 option
+        [[ "$option" != $'\n' ]] && echo
+        case "$option" in
+            [nN]) update_last-updated_file ;;
+            *) update_ohmyzsh ;;
+        esac
     fi
-  else
-    # create the zsh file
-    _update_zsh_update
-  fi
+}
 
-  rmdir $ZSH/log/update.lock
-fi
+unset -f current_epoch update_last-updated_file update_ohmyzsh