1# Helper function for reading input from a TCP connection.
2# Actually, the input doesn't need to be a TCP connection at all, it
3# is simply an input file descriptor.  However, it must be contained
4# in ${tcp_by_fd[$TCP_SESS]}.  This is set by tcp_open, but may be
5# set by hand.  (Note, however, the blocking/timeout behaviour is usually
6# not implemented for reading from regular files.)
7#
8# The default behaviour is simply to read any single available line from
9# the input fd and print it.  If a line is read, it is stored in the
10# parameter $TCP_LINE; this always contains the last line successfully
11# read.  Any chunk of lines read in are stored in the array $tcp_lines;
12# this always contains a complete list of all lines read in by a single
13# execution of this function and hence may be empty.  The fd corresponding
14# to $TCP_LINE is stored in $TCP_LINE_FD (this can be turned into a
15# session by looking up in $tcp_by_fd).
16#
17# Printed lines are preceded by $TCP_PROMPT.  This may contain two
18# percent escapes: %s for the current session, %f for the current file
19# descriptor.  The default is `T[%s]:'.  The prompt is not printed
20# to per-session logs where the source is unambiguous.
21#
22# The function returns 0 if a read succeeded, even if (using -d) a
23# subsequent read failed.
24#
25# The behaviour is modified by the following options.
26#
27# -a     Read from all fds, not just the one given by TCP_SESS.
28#
29# -b	 The first read blocks until some data is available for reading.
30#
31# -d     Drain all pending input; loop until no data is available.
32#
33# -l sess1,sess2,...
34#        Gives a list of sessions to read on.  Equivalent to
35#        -u ${tcp_by_name[sess1]} -u ${tcp_by_name[sess2]} ...
36#	 Multiple -l options also work.
37#
38# -q     Quiet; if $TCP_SESS is not set, just return 1, but don't print
39#        an error message.
40#
41# -s sess
42#        Gives a single session; the option may be repeated.
43#
44# -t TO  On each read (the only read unless -d was also given), time out
45#        if nothing was available after TO seconds (may be floating point).
46#        Otherwise,  the function will return immediately when no data is
47#	 available.
48#
49#        If combined with -b, the function will always wait for the
50#        first data to become available; hence this is not useful unless
51#        -d is specified along with -b, in which case the timeout applies
52#        to data after the first line.
53# -u fd  Read from fd instead of the default session; may be repeated for
54#        multiple sessions.  Can be a comma-separated list, too.
55# -T TO  This sets an overall timeout, again in seconds.
56
57emulate -L zsh
58setopt extendedglob cbases
59# set -x
60
61zmodload -i zsh/mathfunc
62
63local opt drain line quiet block read_fd all sess key val noprint
64local -A read_fds
65read_fds=()
66float timeout timeout_all endtime
67integer stat
68
69while getopts "abdl:qs:t:T:u:" opt; do
70  case $opt in
71    # Read all sessions.
72    (a) all=1
73	;;
74    # Block until we receive something.
75    (b) block=1
76	;;
77    # Drain all pending input.
78    (d) drain=1
79	;;
80    (l) for sess in ${(s.,.)OPTARG}; do
81	  read_fd=${tcp_by_name[$sess]}
82	  if [[ -z $read_fd ]]; then
83	    print "$0: no such session: $sess" >&2
84	    return 1
85	  fi
86	  read_fds[$read_fd]=1
87	done
88	;;
89
90    # Don't print an error message if there is no TCP connection,
91    # just return 1.
92    (q) quiet=1
93	;;
94    # Add a single session to the list
95    (s) read_fd=${tcp_by_name[$OPTARG]}
96        if [[ -z $read_fd ]]; then
97	    print "$0: no such session: $sess" >&2
98	    return 1
99	fi
100	read_fds[$read_fd]=1
101        ;;
102    # Per-read timeout: wait this many seconds before
103    # each read.
104    (t) timeout=$OPTARG
105        [[ -n $TCP_READ_DEBUG ]] && print "Timeout per-operations is $timeout" >&2
106	;;
107    # Overall timeout: return after this many seconds.
108    (T) timeout_all=$OPTARG
109	;;
110    # Read from given fd(s).
111    (u) for read_fd in ${(s.,.)OPTARG}; do
112	  if [[ $read_fd != (0x[[:xdigit:]]##|[[:digit:]]##) ]]; then
113	    print "Bad fd in $OPTARG" >&2
114	    return 1
115	  fi
116	  read_fds[$((read_fd))]=1
117	done
118	;;
119    (*) [[ $opt != \? ]] && print "Unhandled option, complain: $opt" >&2
120	return 1
121       ;;
122  esac
123done
124
125if [[ -n $all ]]; then
126  read_fds=(${(kv)tcp_by_fd})
127elif (( ! $#read_fds )); then
128  if [[ -z $TCP_SESS ]]; then
129    [[ -z $quiet ]] && print "No tcp connection open." >&2
130    return 1
131  elif [[ -z $tcp_by_name[$TCP_SESS] ]]; then
132    print "TCP session $TCP_SESS has gorn!" >&2
133    return 1
134  fi
135  read_fds[$tcp_by_name[$TCP_SESS]]=1
136fi
137
138typeset -ga tcp_lines
139tcp_lines=()
140
141local helper_stat=2 skip tpat reply REPLY
142float newtimeout
143
144if [[ ${(t)SECONDS} != float* ]]; then
145  # If called from another function, don't override
146  typeset -F TCP_SECONDS_START=$SECONDS
147  # Get extra accuracy by making SECONDS floating point locally
148  typeset -F SECONDS
149fi
150
151if (( timeout_all )); then
152  (( endtime = SECONDS + timeout_all ))
153fi
154
155zmodload -i zsh/zselect
156
157if [[ -n $block ]]; then
158  if (( timeout_all )); then
159    # zselect -t uses 100ths of a second
160    zselect -t $(( int(100*timeout_all + 0.5) )) ${(k)read_fds} ||
161      return $helper_stat
162  else
163    zselect ${(k)read_fds} || return $helper_stat
164  fi
165fi
166
167while (( ${#read_fds} )); do
168  if [[ -n $block ]]; then
169    # We already have data waiting this time through.
170    unset block
171  else
172    if (( timeout_all )); then
173      (( (newtimeout = endtime - SECONDS) <= 0 )) && return 2
174      if (( ! timeout || newtimeout < timeout )); then
175	(( timeout = newtimeout ))
176      fi
177    fi
178    if (( timeout )); then
179      if [[ -n $TCP_READ_DEBUG ]]; then
180	print "[tcp_read: selecting timeout $timeout on ${(k)read_fds}]" >&2
181      fi
182      zselect -t $(( int(timeout*100 + 0.5) )) ${(k)read_fds} ||
183        return $helper_stat
184    else
185      if [[ -n $TCP_READ_DEBUG ]]; then
186	print "[tcp_read: selecting no timeout on ${(k)read_fds}]" >&2
187      fi
188      zselect -t 0 ${(k)read_fds} || return $helper_stat
189    fi
190  fi
191  if [[ -n $TCP_READ_DEBUG ]]; then
192    print "[tcp_read: returned fds ${reply}]" >&2
193  fi
194  for read_fd in ${reply[2,-1]}; do
195    if ! read -u $read_fd -r line; then
196      unset "read_fds[$read_fd]"
197      stat=1
198      continue
199    fi
200
201    helper_stat=0
202    sess=${tcp_by_fd[$read_fd]}
203
204    # Handle user-defined triggers
205    noprint=${TCP_SILENT:+-q}
206    if (( ${+tcp_on_read} )); then
207      # Call the function given in the key for each matching value.
208      # It is this way round because function names must be
209      # unique, while patterns do not need to be.  Furthermore,
210      # this keeps the use of subscripting under control.
211      for key val in ${(kv)tcp_on_read}; do
212	if [[ $line = ${~val} ]]; then
213	  $key "$sess" "$line" || noprint=-q
214	fi
215      done
216    fi
217
218    tcp_output -P "${TCP_PROMPT=<-[%s] }" -S $sess -F $read_fd \
219        $noprint -- "$line"
220    # REPLY is now set to the line with an appropriate prompt.
221    tcp_lines+=($REPLY)
222    typeset -g TCP_LINE="$REPLY" TCP_LINE_FD="$read_fd"
223
224    # Only handle one line from one device at a time unless draining.
225    [[ -z $drain ]] && return $stat
226  done
227done
228
229return $stat
230