From 3c777ebf1b24826363bea6bc29afa9ba038b10e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marc=20Cornell=C3=A0?= <marc.cornella@live.com>
Date: Tue, 9 Jun 2020 19:38:08 +0200
Subject: [PATCH] update: refactor and fix logic in check_for_upgrade.sh
 (#8939)

---
 oh-my-zsh.sh               |   5 --
 tools/check_for_upgrade.sh | 132 ++++++++++++++++++++++---------------
 2 files changed, 79 insertions(+), 58 deletions(-)

diff --git a/oh-my-zsh.sh b/oh-my-zsh.sh
index 4025f0025..4b4191c01 100644
--- a/oh-my-zsh.sh
+++ b/oh-my-zsh.sh
@@ -4,11 +4,6 @@ if [[ -z "$ZSH_CACHE_DIR" ]]; then
   ZSH_CACHE_DIR="$ZSH/cache"
 fi
 
-# 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
-
 # Check for updates on initial load...
 if [ "$DISABLE_AUTO_UPDATE" != "true" ]; then
   env ZSH=$ZSH ZSH_CACHE_DIR=$ZSH_CACHE_DIR DISABLE_UPDATE_PROMPT=$DISABLE_UPDATE_PROMPT zsh -f $ZSH/tools/check_for_upgrade.sh
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
-
-zmodload zsh/datetime
-
-function _current_epoch() {
-  echo $(( $EPOCHSECONDS / 60 / 60 / 24 ))
-}
-
-function _update_zsh_update() {
-  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
-}
-
-epoch_target=$UPDATE_ZSH_DAYS
-if [[ -z "$epoch_target" ]]; then
-  # Default to old behavior
-  epoch_target=13
+# 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 upgrade if the current user doesn't have write permissions for the
-# oh-my-zsh directory.
-[[ -w "$ZSH" ]] || return 0
+# 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
 
-# Cancel upgrade if git is unavailable on the system
-whence git >/dev/null || return 0
 
-if mkdir "$ZSH/log/update.lock" 2>/dev/null; then
-  if [ -f ${ZSH_CACHE_DIR}/.zsh-update ]; then
-    . ${ZSH_CACHE_DIR}/.zsh-update
+function current_epoch() {
+    zmodload zsh/datetime
+    echo $(( EPOCHSECONDS / 60 / 60 / 24 ))
+}
 
-    if [[ -z "$LAST_EPOCH" ]]; then
-      _update_zsh_update
-      rmdir $ZSH/log/update.lock # TODO: fix later
-      return 0
-    fi
+function update_last-updated_file() {
+    echo "LAST_EPOCH=$(current_epoch)" >! "${ZSH_CACHE_DIR}/.zsh-update"
+}
 
-    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
+function update_ohmyzsh() {
+    ZSH="$ZSH" sh "$ZSH/tools/upgrade.sh"
+    update_last-updated_file
+}
+
+() {
+    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
     fi
-  else
-    # create the zsh file
-    _update_zsh_update
-  fi
 
-  rmdir $ZSH/log/update.lock
-fi
+    # Check for lock directory
+    if ! command mkdir "$ZSH/log/update.lock" 2>/dev/null; then
+        return
+    fi
+
+    # 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
+
+    # 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
+
+    # 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
+
+    # 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
+}
+
+unset -f current_epoch update_last-updated_file update_ohmyzsh