From c7631e52594858ff18d1ab563e111289f8f8b45e Mon Sep 17 00:00:00 2001
From: Nicholas Yip <nicholasyip@obceit.cc>
Date: Thu, 3 Dec 2020 04:00:28 +0900
Subject: [PATCH] fix(command): Restore environment variables before calling
 `exec`

Previously, the command would overwrite the HOST and PORT variables based on the hostname and port of the service which we would check for being available. This change prevents these variables from being overwritten, which allows you to set them outside the command.

BREAKING CHANGE: HOST, PORT and other internally used environment variables are not overwritten anymore. If you use these, then you need to manually supply them.
---
 wait-for      | 14 +++++++++++---
 wait-for.bats | 13 +++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/wait-for b/wait-for
index e05c733..4ec5263 100755
--- a/wait-for
+++ b/wait-for
@@ -22,6 +22,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+set -- "$@" -- "$TIMEOUT" "$QUIET" "$HOST" "$PORT" "$result"
 TIMEOUT=15
 QUIET=0
 
@@ -52,7 +53,15 @@ wait_for() {
     
     result=$?
     if [ $result -eq 0 ] ; then
-      if [ $# -gt 0 ] ; then
+      if [ $# -gt 6 ] ; then
+        for result in $(seq $(($# - 6))); do
+          result=$1
+          shift
+          set -- "$@" "$result"
+        done
+
+        TIMEOUT=$2 QUIET=$3 HOST=$4 PORT=$5 result=$6
+        shift 6
         exec "$@"
       fi
       exit 0
@@ -69,8 +78,7 @@ wait_for() {
   exit 1
 }
 
-while [ $# -gt 0 ]
-do
+while :; do
   case "$1" in
     *:* )
     HOST=$(printf "%s\n" "$1"| cut -d : -f 1)
diff --git a/wait-for.bats b/wait-for.bats
index e5ff462..d75f0cc 100644
--- a/wait-for.bats
+++ b/wait-for.bats
@@ -32,3 +32,16 @@
   [ "$status" -ne 0 ]
   [ "$output" != "success" ]
 }
+
+@test "environment variables should be restored for command invocation" {
+  HOST=success run ./wait-for -t 1 google.com:80 -- sh -c 'echo "$HOST"'
+
+  [ "$output" = "success" ]
+}
+
+@test "unset environment variables should be restored as unset for command invocation" {
+  run ./wait-for -t 1 google.com:80 -- sh -uc 'echo "$HOST"'
+
+  [ "$status" -ne 0 ]
+  [ "$output" != "google.com" ]
+}