From b4ec9b633b0b3ccbb74dc7d0c029e8ec3d2a33cf Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Fri, 23 Oct 2020 21:38:06 +0300
Subject: [PATCH] Implement `async_job -s` and perform safer zpty writes
 (breaking change)

This change borrows the idea from #45 to introduce newlines to `zpty -w`
as well.

Consequently, when we want to properly escape _all_ inputs to
`async_job`, we can no longer accept scripts as the first argument,
as-is. Furthermore, it was discovered that the current implementation
is flawed in the sense that it's impossible to execute a single command
where the name or path has spaces in it (how did we not notice this
before?).

For this reason, `async_job` received the `-s` argument which allows
running scripts, for example:

	async_job myworker -s 'print hello; print world'

Whereas a command with spaces can now work without changes:

	async_job myworker /path/to/my\ executable
---
 async.zsh      | 50 +++++++++++++++++++++++++++++++++-----------------
 async_test.zsh | 36 +++++++++++++++++++++++++++---------
 2 files changed, 60 insertions(+), 26 deletions(-)

diff --git a/async.zsh b/async.zsh
index 71f3431..e12727c 100644
--- a/async.zsh
+++ b/async.zsh
@@ -226,10 +226,15 @@ _async_worker() {
 			return $(( 127 + 1 ))
 		}
 
-		# We need to clean the input here because sometimes when a zpty
-		# has died and been respawned, messages will be prefixed with a
-		# carraige return (\r, or \C-M).
-		request=${request#$'\C-M'}
+		# Remove CRLF from the request input, any CRLF that are part of
+		# the request are escaped and won't be removed. Previously we
+		# did not escape everything and had to clean the input because
+		# sometimes when a zpty has died and been respawned, messages
+		# will be prefixed with a carraige return (\r, or \C-M).
+		request=${${request//$'\r'}//$'\n'}
+		if [[ -z $request ]]; then
+			continue
+		fi
 
 		# Check for non-job commands sent to worker
 		case $request in
@@ -237,12 +242,17 @@ _async_worker() {
 			_async_eval*) do_eval=1;;
 		esac
 
-		# Parse the request using shell parsing (z) to allow commands
-		# to be parsed from single strings and multi-args alike.
-		cmd=("${(z)request}")
-
-		# Name of the job (first argument).
-		job=$cmd[1]
+		# Check if request is a script argument.
+		if [[ $request == '-s '* ]]; then
+			cmd=("${(Qz)request}")
+			shift cmd
+			job=$cmd[1]  # Legacy, name of first command...
+		else
+			# Parse the request using shell parsing (z) to allow commands
+			# to be parsed from single strings and multi-args alike.
+			cmd=("${(z)request}")
+			job=$cmd[1]
+		fi
 
 		# Check if a worker should perform unique jobs, unless
 		# this is an eval since they run synchronously.
@@ -418,14 +428,19 @@ _async_send_job() {
 		return 1
 	}
 
-	zpty -w $worker "$@"$'\0'
+	# The command is surrounded in newlines to enourage better
+	# behavior from the zpty worker, these are later stripped.
+	zpty -w $worker $'\n'"$@"$'\0\n'
 }
 
 #
 # Start a new asynchronous job on specified worker, assumes the worker is running.
 #
 # usage:
-# 	async_job <worker_name> <my_function> [<function_params>]
+# 	async_job <worker_name> [-s] <my_function> [<function_params>]
+#
+# opts:
+# 	-s interpret the argument as a script
 #
 async_job() {
 	setopt localoptions noshwordsplit noksharrays noposixidentifiers noposixstrings
@@ -434,21 +449,22 @@ async_job() {
 
 	local -a cmd
 	cmd=("$@")
-	if (( $#cmd > 1 )); then
-		cmd=(${(q)cmd})  # Quote special characters in multi argument commands.
-	fi
+	cmd=(${(q)cmd})  # Quote special characters in multi argument commands.
 
 	_async_send_job $0 $worker "$cmd"
 }
 
 #
-# Evaluate a command (like async_job) inside the async worker, then worker environment can be manipulated. For example,
+# Evaluate a command inside the async worker (similar to async_job) so that the worker environment can be manipulated. For example,
 # issuing a cd command will change the PWD of the worker which will then be inherited by all future async jobs.
 #
 # Output will be returned via callback, job name will be [async/eval].
 #
 # usage:
-# 	async_worker_eval <worker_name> <my_function> [<function_params>]
+# 	async_worker_eval <worker_name> [-s] <my_function> [<function_params>]
+#
+# opts:
+# 	-s interpret the argument as a script
 #
 async_worker_eval() {
 	setopt localoptions noshwordsplit noksharrays noposixidentifiers noposixstrings
diff --git a/async_test.zsh b/async_test.zsh
index 143aa34..e016b88 100644
--- a/async_test.zsh
+++ b/async_test.zsh
@@ -143,7 +143,7 @@ test_async_process_results_stress() {
 
 	integer iter=40 timeout=5
 	for i in {1..$iter}; do
-		async_job test "print -n $i"
+		async_job test -s "print -n $i"
 	done
 
 	float start=$EPOCHSECONDS
@@ -174,7 +174,7 @@ test_async_process_results_stress() {
 	# Test with longer running commands (sleep, then print).
 	iter=40
 	for i in {1..$iter}; do
-		async_job test "sleep 1 && print -n $i"
+		async_job test -s "sleep 1 && print -n $i"
 		sleep 0.00001
 		(( iter % 6 == 0 )) && async_process_results test cb
 	done
@@ -210,7 +210,7 @@ test_async_job_multiple_commands_in_multiline_string() {
 
 	async_start_worker test
 	# Test multi-line (single string) command.
-	async_job test 'print "hi\n  123 "'$'\nprint -n bye'
+	async_job test -s 'print "hi\n  123 "'$'\nprint -n bye'
 	while ! async_process_results test cb; do :; done
 	async_stop_worker test
 
@@ -219,6 +219,24 @@ test_async_job_multiple_commands_in_multiline_string() {
 	[[ $result[3] = $want ]] || t_error "want output: ${(Vq-)want}, got ${(Vq-)result[3]}"
 }
 
+test_async_job_command_with_space() {
+	local -a result
+	cb() { result=("$@") }
+
+	async_start_worker test
+	# Test multi-line (single string) command.
+	async_job test '/abra cadabra'
+	while ! async_process_results test cb; do :; done
+	async_stop_worker test
+
+	local want='/abra\ cadabra'
+	[[ $result[1] = $want ]] || t_error "want command name: ${(Vq-)want}, got ${(Vq-)result[1]}"
+
+	# Check that we get a command not found.
+	want='/abra cadabra'
+	[[ $result[5] = *$want ]] || t_error "want stderr to contain: ${(Vq-)want}, got ${(Vq-)result[5]}"
+}
+
 test_async_job_git_status() {
 	local -a result
 	cb() { result=("$@") }
@@ -322,7 +340,7 @@ test_async_worker_notify_sigwinch() {
 	async_start_worker test -n
 	async_register_callback test cb
 
-	async_job test 'sleep 0.1; print hi'
+	async_job test -s 'sleep 0.1; print hi'
 
 	while (( ! $#result )); do sleep 0.01; done
 
@@ -437,9 +455,9 @@ test_async_worker_update_pwd() {
 	async_start_worker test1
 	t_defer async_stop_worker test1
 
-	async_job test1 'print $PWD'
+	async_job test1 -s 'print $PWD'
 	async_worker_eval test1 'print -n foo; cd ..; print -n bar; print -n -u2 baz'
-	async_job test1 'print $PWD'
+	async_job test1 -s 'print $PWD'
 
 	start=$EPOCHREALTIME
 	while (( EPOCHREALTIME - start < 2.0 && $#result < 2 )); do
@@ -468,9 +486,9 @@ test_async_worker_update_pwd_and_env() {
 	async_start_worker test1
 	t_defer async_stop_worker test1
 
-	async_job test1 "print -n $myenv"
+	async_job test1 -s "print -n $myenv"
 	async_worker_eval test1 "cd ..; export myenv=${(q)input}"
-	async_job test1 'print -n $myenv'
+	async_job test1 -s 'print -n $myenv'
 
 	start=$EPOCHREALTIME
 	while (( EPOCHREALTIME - start < 2.0 && $#result < 2 )); do
@@ -623,7 +641,7 @@ test_zle_watcher() {
 		t_fatal "want _async_zle_watcher to be registered as zle watcher, got output ${(Vq-)result}"
 	}
 
-	zpty_run async_job test 'print hello world' || t_fatal "could not send async_job command"
+	zpty_run async_job test -s 'print hello world' || t_fatal "could not send async_job command"
 
 	zpty -r -m zsh result "*print 0 'hello world'*" || {
 		t_fatal "want \"print 0 'hello world'\", got output ${(Vq-)result}"