18810: Various enhancements and bug fixes for the TCP function suite

9 files changed, 130 insertions, 16 deletions

diff --git a/Functions/TCP/tcp_expect b/Functions/TCP/tcp_expect
index 14963a3e6..e49f93804 100644
--- a/Functions/TCP/tcp_expect
+++ b/Functions/TCP/tcp_expect
@@ -37,8 +37,12 @@
 emulate -L zsh
 setopt extendedglob
 
-# Get extra accuracy by making SECONDS floating point locally
-typeset -F SECONDS
+if [[ ${(t)SECONDS} != float* ]]; then
+    # If called from another function, use that
+    typeset -F TCP_SECONDS_START=$SECONDS
+    # Get extra accuracy by making SECONDS floating point locally
+    typeset -F SECONDS
+fi
 
 # Variables are all named _expect_* to avoid problems with the -p param.
 local _expect_opt _expect_pvar
diff --git a/Functions/TCP/tcp_open b/Functions/TCP/tcp_open
index f762d932a..1f7a65c9a 100644
--- a/Functions/TCP/tcp_open
+++ b/Functions/TCP/tcp_open
@@ -53,11 +53,19 @@
 emulate -L zsh
 setopt extendedglob cbases
 
+# Global set up for TCP function suite.
+
 zmodload -i zsh/net/tcp || return 1
 zmodload -i zsh/zutil
 autoload -U tcp_alias tcp_close tcp_command tcp_expect tcp_fd_handler
 autoload -U tcp_log tcp_output tcp_proxy tcp_read tcp_rename tcp_send
-autoload -U tcp_sess tcp_spam tcp_talk tcp_wait
+autoload -U tcp_sess tcp_spam tcp_talk tcp_wait tcp_point tcp_shoot
+
+# TCP_SECONDS_START is only set if we override TCP_SECONDS locally,
+# so provide a global value for convenience.  Should probably always be 0.
+(( ${+TCP_SECONDS_START} )) || typeset -gF TCP_SECONDS_START
+
+# Processing for new connection.
 
 local opt accept fake nozle sessfile sess quiet
 local -a sessnames sessargs
diff --git a/Functions/TCP/tcp_output b/Functions/TCP/tcp_output
index 69177bae0..781c46c33 100644
--- a/Functions/TCP/tcp_output
+++ b/Functions/TCP/tcp_output
@@ -1,7 +1,7 @@
 emulate -L zsh
 setopt extendedglob
 
-local opt tprompt sess read_fd tpat quiet
+local opt tprompt sess read_fd tpat quiet cursess
 
 while getopts "F:P:qS:" opt; do
   case $opt in
@@ -29,7 +29,12 @@ fi
 # where data is coming from; also, it allows more predictable
 # behaviour in tcp_expect.
 if [[ -n $tprompt ]]; then
-  zformat -f REPLY $tprompt "s:$sess" "f:$read_fd"
+  if [[ $sess = $TCP_SESS ]]; then
+      cursess="c:1"
+  else
+      cursess="c:0"
+  fi
+  zformat -f REPLY $tprompt "s:$sess" "f:$read_fd" $cursess
   # We will pass this back up.
   REPLY="$REPLY$*"
 else
diff --git a/Functions/TCP/tcp_point b/Functions/TCP/tcp_point
new file mode 100644
index 000000000..6a8f75d3a
--- /dev/null
+++ b/Functions/TCP/tcp_point
@@ -0,0 +1,29 @@
+emulate -L zsh
+setopt extendedglob cbases
+
+
+if [[ $# -ne 1 ]]; then
+    print "Usage: $0 port
+Listen on the given port; send anything that arrives to standard output." >&2
+    return 1
+fi
+
+local REPLY lfd afd
+if ! ztcp -l $1; then
+    print "Failed to listen on port $1" >&2
+    return 1
+fi
+
+lfd=$REPLY
+
+if ! ztcp -a $lfd; then
+    print "Failed to accept on fd $lfd" >&2
+    ztcp -c $lfd
+fi
+
+afd=$REPLY
+
+cat <&$afd
+
+ztcp -c $lfd
+ztcp -c $afd
diff --git a/Functions/TCP/tcp_read b/Functions/TCP/tcp_read
index d9aa47e0e..51e5356cc 100644
--- a/Functions/TCP/tcp_read
+++ b/Functions/TCP/tcp_read
@@ -60,7 +60,7 @@ setopt extendedglob cbases
 
 zmodload -i zsh/mathfunc
 
-local opt drain line quiet block read_fd all sess
+local opt drain line quiet block read_fd all sess key val
 local -A read_fds
 read_fds=()
 float timeout timeout_all endtime
@@ -139,8 +139,12 @@ tcp_lines=()
 local helper_stat=2 skip tpat reply REPLY
 float newtimeout
 
-# Get extra accuracy by making SECONDS floating point locally
-typeset -F SECONDS
+if [[ ${(t)SECONDS} != float* ]]; then
+    # If called from another function, don't override
+    typeset -F TCP_SECONDS_START=$SECONDS
+    # Get extra accuracy by making SECONDS floating point locally
+    typeset -F SECONDS
+fi
 
 if (( timeout_all )); then
   (( endtime = SECONDS + timeout_all ))
@@ -194,11 +198,23 @@ while (( ${#read_fds} )); do
 
     helper_stat=0
     sess=${tcp_by_fd[$read_fd]}
-    tcp_output -P "${TCP_PROMPT:=<-[%s] }" -S $sess -F $read_fd \
+    tcp_output -P "${TCP_PROMPT=<-[%s] }" -S $sess -F $read_fd \
       ${TCP_SILENT:+-q} "$line"
     # REPLY is now set to the line with an appropriate prompt.
     tcp_lines+=($REPLY)
     TCP_LINE=$REPLY TCP_LINE_FD=$read_fd
+
+    # Handle user-defined triggers
+    if (( ${+tcp_on_read} )); then
+	# Call the function given in the key for each matching value.
+	# It is this way round because function names must be
+	# unique, while patterns do not need to be.  Furthermore,
+	# this keeps the use of subscripting under control.
+	for key val in ${(kv)tcp_on_read}; do
+	    [[ $line = ${~val} ]] && $key "$sess" "$line"
+	done
+    fi
+
     # Only handle one line from one device at a time unless draining.
     [[ -z $drain ]] && return $stat
   done
diff --git a/Functions/TCP/tcp_send b/Functions/TCP/tcp_send
index e7dfca771..c5f902f71 100644
--- a/Functions/TCP/tcp_send
+++ b/Functions/TCP/tcp_send
@@ -3,6 +3,7 @@ setopt extendedglob cbases
 
 local opt quiet all sess fd nonewline
 local -a sessions write_fds
+integer mystat
 
 while getopts "al:nqs:" opt; do
     case $opt in
@@ -56,6 +57,11 @@ local TCP_SESS
 
 for TCP_SESS in $sessions; do
     fd=${tcp_by_name[$TCP_SESS]}
+    if [[ -z $fd ]]; then
+	print "No such session: $TCP_SESS" >&2
+	mystat=1
+	continue
+    fi
     print $nonewline -r -- $* >&$fd
     if [[ $? -ne 0 || -n $TCP_FD_CLOSED ]]; then
 	print "Session ${TCP_SESS}: fd $fd unusable." >&2
@@ -65,3 +71,5 @@ for TCP_SESS in $sessions; do
 	tcp_output -P "$TCP_OUTPUT" -S $TCP_SESS -F $fd -q "${(j. .)*}"
     fi
 done
+
+return $mystat
diff --git a/Functions/TCP/tcp_shoot b/Functions/TCP/tcp_shoot
new file mode 100644
index 000000000..8ff9866ba
--- /dev/null
+++ b/Functions/TCP/tcp_shoot
@@ -0,0 +1,21 @@
+emulate -L zsh
+setopt extendedglob
+
+local REPLY tfd
+
+if [[ $# -ne 2 ]]; then
+    print "Usage: tcp_dump host port
+Connect to the given host and port; send standard input.">&2
+    return 1
+fi
+
+if ! ztcp $1 $2; then
+    print "Failed to open connection to host $1 port $2" >&2
+    return 1
+fi
+
+tfd=$REPLY
+
+cat >&$tfd
+
+ztcp -c $tfd
diff --git a/Functions/TCP/tcp_spam b/Functions/TCP/tcp_spam
index 8b84cd0d2..345be7f21 100644
--- a/Functions/TCP/tcp_spam
+++ b/Functions/TCP/tcp_spam
@@ -4,6 +4,8 @@
 #    If not given and tcp_spam_list is set to a list of sessions,
 #    only those will be spammed.  If tcp_no_spam_list is set, those
 #    will (also) be excluded from spamming.
+# -e use `eval' to run the command list instead of executing as
+#    a normal command line.
 # -l sess1,sess2    give comma separated list of sessions to spam
 # -r reverse, spam in opposite order (default is alphabetic, -r means
 #    omegapsiic).  Note tcp_spam_list is not sorted (but may be reversed).
@@ -19,14 +21,17 @@
 emulate -L zsh
 setopt extendedglob
 
-local TCP_SESS cmd opt verbose reverse sesslist transmit all
+local cursess=$TCP_SESS sessstr
+local TCP_SESS cmd opt verbose reverse sesslist transmit all eval
 local match mbegin mend REPLY
 local -a sessions
 
-while getopts "al:rtv" opt; do
+while getopts "ael:rtv" opt; do
     case $opt in
 	(a) all=1
 	    ;;
+	(e) eval=1
+	    ;;
 	(l) sessions+=(${(s.,.)OPTARG})
 	    ;;
 	(r) reverse=1
@@ -82,7 +87,7 @@ fi
 
 if [[ -n $transmit ]]; then
   cmd=tcp_send
-else
+elif [[ -z $eval ]]; then
   cmd=$1
   shift
 fi
@@ -95,7 +100,18 @@ for TCP_SESS in $sessions; do
     tcp_on_spam $TCP_SESS $cmd $*
     [[ $REPLY = done ]] && continue
   fi
-  [[ -n $verbose ]] && zformat -f REPLY $TCP_PROMPT "s:$TCP_SESS" \
-    "f:${tcp_by_name[$TCP_SESS]}" && print -r $REPLY
-  eval $cmd '$*'
+  if [[ -n $verbose ]]; then
+      if [[ $TCP_SESS = $cursess ]]; then
+	  sessstr="c:1"
+      else
+	  sessstr="c:0"
+      fi
+      zformat -f REPLY $TCP_PROMPT "s:$TCP_SESS" \
+	  "f:${tcp_by_name[$TCP_SESS]}" $sessstr && print -r $REPLY
+  fi
+  if [[ -n $eval ]]; then
+      eval $*
+  else
+      eval $cmd '$*'
+  fi
 done
diff --git a/Functions/TCP/tcp_wait b/Functions/TCP/tcp_wait
index d18068a66..ee04cf23b 100644
--- a/Functions/TCP/tcp_wait
+++ b/Functions/TCP/tcp_wait
@@ -1,7 +1,14 @@
 # Wait for given number of seconds, reading any data from
 # all TCP connections while doing so.
 
-typeset -F SECONDS to end
+if [[ ${(t)SECONDS} != float* ]]; then
+    # If called from tcp_expect, don't override
+    typeset -F TCP_SECONDS_START=$SECONDS
+    # Get extra accuracy by making SECONDS floating point locally
+    typeset -F SECONDS
+fi
+
+typeset to end
 
 (( to = $1, end = SECONDS + to ))
 while (( SECONDS < end )); do