Prevent infinite loop if a shim is linked and called from a nonstandard location with "system" active

libexec/pyenv-exec

@@ -29,6 +29,14 @@ if [ -z "$PYENV_COMMAND" ]; then
   exit 1
 fi
 
+if [[ -n $_PYENV_SHIM_PATH ]]; then
+  PROGRAM="$(echo "$PYENV_COMMAND" | tr a-z- A-Z_ | sed 's/[^A-Z0-9_]/_/g')"
+  NR_PYENV_SHIM_PATHS_PROGRAM="_PYENV_SHIM_PATHS_${PROGRAM}"
+  export -- "$NR_PYENV_SHIM_PATHS_PROGRAM=$_PYENV_SHIM_PATH${!NR_PYENV_SHIM_PATHS_PROGRAM:+:${!NR_PYENV_SHIM_PATHS_PROGRAM}}"
+  unset PROGRAM NR_PYENV_SHIM_PATHS_PROGRAM
+  unset _PYENV_SHIM_PATH
+fi
+
 PYENV_COMMAND_PATH="$(pyenv-which "$PYENV_COMMAND")"
 PYENV_BIN_PATH="${PYENV_COMMAND_PATH%/*}"
 export PYENV_VERSION

libexec/pyenv-rehash

@@ -82,6 +82,10 @@ set -e
 program="\${0##*/}"
 
 export PYENV_ROOT="$PYENV_ROOT"
+SHIM_PATH=\${0%/*}
+if [[ \$SHIM_PATH != "$PYENV_ROOT/shims" ]]; then
+  export _PYENV_SHIM_PATH="\$SHIM_PATH"
+fi
 exec "$(command -v pyenv)" exec "\$program" "\$@"
 SH
   chmod +x "$PROTOTYPE_SHIM_PATH"

libexec/pyenv-which

@@ -42,15 +42,19 @@ done
 
 
 remove_from_path() {
-  local path_to_remove="$1"
+  local -a paths_to_remove
-  local path_before
+  IFS=: paths_to_remove=($1)
+  local path_to_remove path_before
   local result=":${PATH//\~/$HOME}:"
-  while [ "$path_before" != "$result" ]; do
+  for path_to_remove in "${paths_to_remove[@]}"; do
-    path_before="$result"
+    while true; do
-    result="${result//:$path_to_remove:/:}"
+      path_before="$result"
+      result="${result//:$path_to_remove:/:}"
+      if [[ ${#path_before} == "${#result}" ]]; then break; fi
+    done
   done
-  result="${result%:}"
+  result="${result:1:${#result}-2}"
-  echo "${result#:}"
+  echo "$result"
 }
 
 if [ -z "$PYENV_COMMAND" ]; then
@@ -66,8 +70,10 @@ declare -a nonexistent_versions
 
 for version in "${versions[@]}" "$system"; do
   if [ "$version" = "system" ]; then
-    PATH="$(remove_from_path "${PYENV_ROOT}/shims")"
+    PROGRAM="$(echo "$PYENV_COMMAND" | tr a-z- A-Z_ | sed 's/[^A-Z0-9_]/_/g')"
-    PYENV_COMMAND_PATH="$(command -v "$PYENV_COMMAND" || true)"
+    NR_CUSTOM_SHIM_PATHS="_PYENV_SHIM_PATHS_$PROGRAM"
+    SEARCH_PATH="$(remove_from_path "${PYENV_ROOT}/shims${!NR_CUSTOM_SHIM_PATHS:+:${!NR_CUSTOM_SHIM_PATHS}}")"
+    PYENV_COMMAND_PATH="$(PATH="$SEARCH_PATH" command -v "$PYENV_COMMAND" || true)"
   else
     # $version may be a prefix to be resolved by pyenv-latest
     version_path="$(pyenv-prefix "${version}" 2>/dev/null)" || \

test/exec.bats

@@ -116,3 +116,26 @@ OUT
   PYENV_VERSION=system:custom run pyenv-exec python3
   assert_success "$PATH"
 }
+
+@test "sets/adds to _PYENV_SHIM_PATHS_{PROGRAM} when _PYENV_SHIM_PATH is set, unsets _PYENV_SHIM_PATH" {
+  progname='123;wacky-prog.name ^%$#'
+  envvarname="_PYENV_SHIM_PATHS_123_WACKY_PROG_NAME_____"
+  create_path_executable "$progname" <<!
+echo $envvarname="\$$envvarname"
+echo _PYENV_SHIM_PATH="\$_PYENV_SHIM_PATH"
+!
+  _PYENV_SHIM_PATH=/unusual/shim/location run pyenv-exec "$progname"
+  assert_success
+  assert_output <<!
+$envvarname=/unusual/shim/location
+_PYENV_SHIM_PATH=
+!
+
+  eval "export ${envvarname}=/another/shim/location"
+  _PYENV_SHIM_PATH=/unusual/shim/location run pyenv-exec "$progname"
+  assert_success
+  assert_output <<!
+$envvarname=/unusual/shim/location:/another/shim/location
+_PYENV_SHIM_PATH=
+!
+}

test/rehash.bats

@@ -127,3 +127,28 @@ hash -r 2>/dev/null || true"
   assert_success
   assert [ -x "${PYENV_ROOT}/shims/python" ]
 }
+
+@test "shim sets _PYENV_SHIM_PATH when linked from elsewhere" {
+  export PYENV_VERSION="custom"
+  create_alt_executable python3
+  #must stub pyenv before rehash 'cuz the path is hardcoded into shims
+  create_stub pyenv <<!
+[[ \$1 == 'exec' ]] && \
+echo _PYENV_SHIM_PATH="\$_PYENV_SHIM_PATH"
+!
+  pyenv-rehash
+  mkdir -p "${PYENV_TEST_DIR}/alt-shim"
+  ln -s "${PYENV_ROOT}/shims/python3" "${PYENV_TEST_DIR}/alt-shim/python3"
+  
+  run "${PYENV_TEST_DIR}/alt-shim/python3"
+  assert_success
+  assert_output <<!
+_PYENV_SHIM_PATH=${PYENV_TEST_DIR}/alt-shim
+!
+
+  run python3
+  assert_success
+  assert_output <<!
+_PYENV_SHIM_PATH=
+!
+}

test/shims-linked-from-elsewhere.bats

@@ -0,0 +1,39 @@
+#!/usr/bin/env bats
+
+load test_helper
+
+@test "shims linked from elsewhere are chained, other programs at the same paths are unaffected (integration)" {
+  bats_require_minimum_version 1.5.0
+  progname=shimmed_program
+  called_progname=another_program
+  create_alt_executable_in_version "custom" "$progname" <<!
+echo "called from \$0"
+pyenv-exec "$called_progname"
+!
+  pyenv-rehash
+  PATH="$(path_without "$progname" "$called_progname")"
+  for disguised_shim_path in "$BATS_TEST_TMPDIR/"{weird-location,even/weirder/location}; do
+    mkdir -p "$disguised_shim_path"
+    ln -s "${PYENV_ROOT}/shims/$progname" "$disguised_shim_path/$progname"
+    PATH="$PATH:$disguised_shim_path"
+  done
+  create_executable "$disguised_shim_path" "$called_progname" <<!
+echo "convoluted call success!"
+!
+  real_path="$BATS_TEST_TMPDIR/real-location"
+  mkdir -p "$real_path"
+  ln -s "${PYENV_ROOT}/versions/custom/bin/$progname" "$real_path/$progname"
+  PATH="$PATH:$real_path"
+  
+  run "$progname"
+  assert_success
+  assert_output <<!
+called from $real_path/$progname
+convoluted call success!
+!
+
+  rm "$real_path/$progname"
+  run -127 "$progname"
+  assert_failure
+  assert_line 0 "pyenv: shimmed_program: command not found"
+}

test/test_helper.bash

@@ -28,7 +28,7 @@ setup() {
   PATH="${BATS_TEST_DIRNAME%/*}/libexec:$PATH"
   PATH="${BATS_TEST_DIRNAME}/libexec:$PATH"
   PATH="${PYENV_ROOT}/shims:$PATH"
-  export PATH
+  PATH="${BATS_TEST_TMPDIR}/stubs:$PATH"
 
   for xdg_var in `env 2>/dev/null | grep ^XDG_ | cut -d= -f1`; do unset "$xdg_var"; done
   unset xdg_var
@@ -128,7 +128,7 @@ path_without() {
       if [ "$found" != "${PYENV_ROOT}/shims" ]; then
         alt="${PYENV_TEST_DIR}/$(echo "${found#/}" | tr '/' '-')"
         mkdir -p "$alt"
-        for util in bash head cut readlink greadlink; do
+        for util in bash head cut readlink greadlink tr sed; do
           if [ -x "${found}/$util" ]; then
             ln -s "${found}/$util" "${alt}/$util"
           fi
@@ -156,9 +156,13 @@ create_alt_executable_in_version() {
   create_executable "${PYENV_ROOT}/versions/$version/bin" "$@"
 }
 
+create_stub() {
+  create_executable "${BATS_TEST_TMPDIR}/stubs" "$@"
+}
+
 create_executable() {
-  bin="${1:?}"
+  local bin="${1:?}"
-  name="${2:?}"
+  local name="${2:?}"
   shift 2
   mkdir -p "$bin"
   local payload

test/which.bats

@@ -193,3 +193,29 @@ exit
 pyenv: py.test: command not found
 OUT
 }
+
+@test "excludes paths in _PYENV_SHIM_PATHS_{PROGRAM} from search only for PROGRAM" {
+  progname='123;wacky-prog.name ^%$#'
+  envvarname="_PYENV_SHIM_PATHS_123_WACKY_PROG_NAME_____"
+  create_path_executable "$progname"
+
+  for dir_ in "$BATS_TEST_TMPDIR/alt-path"{1,2}; do
+    mkdir -p "$dir_"
+    ln -s "${PYENV_TEST_DIR}/bin/$progname" "$dir_/$progname"
+    eval 'export '"$envvarname"'="$dir_${'"$envvarname"':+:$'"$envvarname"'}"'
+    PATH="$dir_:$PATH"
+  done
+  create_executable "$dir_" "normal_program"
+  
+  run pyenv-which "$progname"
+  assert_success
+  assert_output <<!
+$PYENV_TEST_DIR/bin/$progname
+!
+
+  run pyenv-which "normal_program"
+  assert_success
+  assert_output <<!
+$dir_/normal_program
+!
+}