about summary refs log tree commit diff
path: root/Functions/TCP/tcp_read
blob: 97da8bf21d017fedc713277f6c7646935a974388 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
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