18202: New TCP function system plus small error message change in ztcp.

1 files changed, 207 insertions, 0 deletions

diff --git a/Functions/TCP/tcp_read b/Functions/TCP/tcp_read
new file mode 100644
index 000000000..97da8bf21
--- /dev/null
+++ b/Functions/TCP/tcp_read
@@ -0,0 +1,207 @@
+# Helper function for reading input from a TCP connection.
+# Actually, the input doesn't need to be a TCP connection at all, it
+# is simply an input file descriptor.  However, it must be contained
+# in ${tcp_by_fd[$TCP_SESS]}.  This is set set by tcp_open, but may be
+# set by hand.  (Note, however, the blocking/timeout behaviour is usually
+# not implemented for reading from regular files.)
+#
+# The default behaviour is simply to read any single available line from
+# the input fd and print it.  If a line is read, it is stored in the
+# parameter $TCP_LINE; this always contains the last line successfully
+# read.  Any chunk of lines read in are stored in the array $tcp_lines;
+# this always contains a complete list of all lines read in by a single
+# execution of this function and hence may be empty.  The fd corresponding
+# to $TCP_LINE is stored in $TCP_LINE_FD (this can be turned into a
+# session by looking up in $tcp_by_fd).
+#
+# Printed lines are preceded by $TCP_PROMPT.  This may contain two
+# percent escapes: %s for the current session, %f for the current file
+# descriptor.  The default is `T[%s]:'.  The prompt is not printed
+# to per-session logs where the source is unambiguous.
+#
+# The function returns 0 if a read succeeded, even if (using -d) a
+# subsequent read failed.
+#
+# The behaviour is modified by the following options.
+#
+# -a     Read from all fds, not just the one given by TCP_SESS.
+#
+# -b	 The first read blocks until some data is available for reading.
+#
+# -d     Drain all pending input; loop until no data is available.
+#
+# -l sess1,sess2,...
+#        Gives a list of sessions to read on.  Equivalent to
+#        -u ${tcp_by_name[sess1]} -u ${tcp_by_name[sess2]} ...
+#	 Multiple -l options also work.
+#
+# -q     Quiet; if $TCP_SESS is not set, just return 1, but don't print
+#        an error message.
+#
+# -s sess
+#        Gives a single session; the option may be repeated.
+#
+# -t TO  On each read (the only read unless -d was also given), time out
+#        if nothing was available after TO seconds (may be floating point).
+#        Otherwise,  the function will return immediately when no data is
+#	 available.
+#
+#        If combined with -b, the function will always wait for the
+#        first data to become available; hence this is not useful unless
+#        -d is specified along with -b, in which case the timeout applies
+#        to data after the first line.
+# -u fd  Read from fd instead of the default session; may be repeated for
+#        multiple sessions.  Can be a comma-separated list, too.
+# -T TO  This sets an overall timeout, again in seconds.
+
+emulate -L zsh
+setopt extendedglob cbases
+# set -x
+
+zmodload -i zsh/mathfunc
+
+local opt drain line quiet block read_fd all sess
+local -A read_fds
+read_fds=()
+float timeout timeout_all endtime
+integer stat
+
+while getopts "abdl:qs:t:T:u:" opt; do
+  case $opt in
+    # Read all sessions.
+    (a) all=1
+	;;
+    # Block until we receive something.
+    (b) block=1
+	;;
+    # Drain all pending input.
+    (d) drain=1
+	;;
+    (l) for sess in ${(s.,.)OPTARG}; do
+	  read_fd=${tcp_by_name[$sess]}
+	  if [[ -z $read_fd ]]; then
+	    print "$0: no such session: $sess" >&2
+	    return 1
+	  fi
+	  read_fds[$read_fd]=1
+	done
+	;;
+
+    # Don't print an error mesage if there is no TCP connection,
+    # just return 1.
+    (q) quiet=1
+	;;
+    # Add a single session to the list
+    (s) read_fd=${tcp_by_name[$OPTARG]}
+        if [[ -z $read_fd ]]; then
+	    print "$0: no such session: $sess" >&2
+	    return 1
+	fi
+	read_fds[$read_fd]=1
+        ;;
+    # Per-read timeout: wait this many seconds before
+    # each read.
+    (t) timeout=$OPTARG
+        [[ -n $TCP_READ_DEBUG ]] && print "Timeout per-operations is $timeout" >&2
+	;;
+    # Overall timeout: return after this many seconds.
+    (T) timeout_all=$OPTARG
+	;;
+    # Read from given fd(s).
+    (u) for read_fd in ${(s.,.)OPTARG}; do
+	  if [[ $read_fd != (0x[[:xdigit:]]##|[[:digit:]]##) ]]; then
+	    print "Bad fd in $OPTARG" >&2
+	    return 1
+	  fi
+	  read_fds[$((read_fd))]=1
+	done
+	;;
+    (*) [[ $opt != \? ]] && print Unhandled option, complain: $opt >&2
+	return 1
+       ;;
+  esac
+done
+
+if [[ -n $all ]]; then
+  read_fds=(${(kv)tcp_by_fd})
+elif (( ! $#read_fds )); then
+  if [[ -z $TCP_SESS ]]; then
+    [[ -z $quiet ]] && print "No tcp connection open." >&2
+    return 1
+  elif [[ -z $tcp_by_name[$TCP_SESS] ]]; then
+    print "TCP session $TCP_SESS has gorn!" >&2
+    return 1
+  fi
+  read_fds[$tcp_by_name[$TCP_SESS]]=1
+fi
+
+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 (( timeout_all )); then
+  (( endtime = SECONDS + timeout_all ))
+fi
+
+zmodload -i zsh/zselect
+
+if [[ -n $block ]]; then
+  if (( timeout_all )); then
+    # zselect -t uses 100ths of a second
+    zselect -t $(( int(100*timeout_all + 0.5) )) ${(k)read_fds} || 
+      return $helper_stat
+  else
+    zselect ${(k)read_fds} || return $helper_stat
+  fi
+fi
+
+while (( ${#read_fds} )); do
+  if [[ -n $block ]]; then
+    # We already have data waiting this time through.
+    unset block
+  else
+    if (( timeout_all )); then
+      (( (newtimeout = endtime - SECONDS) <= 0 )) && return 2
+      if (( ! timeout || newtimeout < timeout )); then
+	(( timeout = newtimeout ))
+      fi
+    fi
+    if (( timeout )); then
+      if [[ -n $TCP_READ_DEBUG ]]; then
+	print "[tcp_read: selecting timeout $timeout on ${(k)read_fds}]" >&2
+      fi
+      zselect -t $(( int(timeout*100 + 0.5) )) ${(k)read_fds} ||
+        return $helper_stat
+    else
+      if [[ -n $TCP_READ_DEBUG ]]; then
+	print "[tcp_read: selecting no timeout on ${(k)read_fds}]" >&2
+      fi
+      zselect -t 0 ${(k)read_fds} || return $helper_stat
+    fi
+  fi
+  if [[ -n $TCP_READ_DEBUG ]]; then
+    print "[tcp_read: returned fds ${reply}]" >&2
+  fi
+  for read_fd in ${reply[2,-1]}; do
+    if ! read -r line <&$read_fd; then
+      unset "read_fds[$read_fd]"
+      stat=1
+      continue
+    fi
+
+    helper_stat=0
+    sess=${tcp_by_fd[$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
+    # Only handle one line from one device at a time unless draining.
+    [[ -z $drain ]] && return $stat
+  done
+done
+
+return $stat