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" ]
+}