3 files changed, 78 insertions, 51 deletions
diff --git a/tools/check_for_upgrade.sh b/tools/check_for_upgrade.sh
index c8dedcf77..c91431c25 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
+            [yY$'\n']) update_ohmyzsh ;;
+            [nN]) update_last_updated_file ;;
+        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
diff --git a/tools/install.sh b/tools/install.sh
index 08f3db1e3..8a93708ca 100755
--- a/tools/install.sh
+++ b/tools/install.sh
@@ -269,11 +269,11 @@ main() {
 		                        /____/                       ....is now installed!
 
 
-		Please look over the ~/.zshrc file to select plugins, themes, and options.
+		Before you scream Oh My Zsh! please look over the ~/.zshrc file to select plugins, themes, and options.
 
-		p.s. Follow us on https://twitter.com/ohmyzsh
-
-		p.p.s. Get stickers, shirts, and coffee mugs at https://shop.planetargon.com/collections/oh-my-zsh
+		• Follow us on Twitter: https://twitter.com/ohmyzsh
+		• Join our Discord server: https://discord.gg/ohmyzsh
+		• Get stickers, shirts, coffee mugs and other swag: https://shop.planetargon.com/collections/oh-my-zsh
 
 	EOF
 	printf "$RESET"
diff --git a/tools/upgrade.sh b/tools/upgrade.sh
index 19ac88c73..de7a8419b 100644
--- a/tools/upgrade.sh
+++ b/tools/upgrade.sh
@@ -62,7 +62,8 @@ then
   printf '%s\____/%s_/ /_/ %s /_/ /_/ /_/%s\__, / %s   /___/%s____/%s_/ /_/  %s\n' $RB_RED $RB_ORANGE $RB_YELLOW $RB_GREEN $RB_BLUE $RB_INDIGO $RB_VIOLET $RB_RESET
   printf '%s    %s        %s           %s /____/ %s       %s     %s          %s\n' $RB_RED $RB_ORANGE $RB_YELLOW $RB_GREEN $RB_BLUE $RB_INDIGO $RB_VIOLET $RB_RESET
   printf "${BLUE}%s\n" "Hooray! Oh My Zsh has been updated and/or is at the current version."
-  printf "${BLUE}${BOLD}%s${RESET}\n" "To keep up on the latest news and updates, follow us on twitter: https://twitter.com/ohmyzsh"
+  printf "${BLUE}${BOLD}%s${RESET}\n" "To keep up on the latest news and updates, follow us on Twitter: https://twitter.com/ohmyzsh"
+  printf "${BLUE}${BOLD}%s${RESET}\n" "Want to get involved in the community? Join our Discord: https://discord.gg/ohmyzsh"
   printf "${BLUE}${BOLD}%s${RESET}\n" "Get your Oh My Zsh swag at: https://shop.planetargon.com/collections/oh-my-zsh"
 else
   printf "${RED}%s${RESET}\n" 'There was an error updating. Try again later?'