Factor out common mainloop code between client and server.

Also improve the socket message output a bit.
This commit is contained in:
Avery Pennarun 2010-10-01 17:36:09 -07:00
parent b0f061e204
commit 84376284db
3 changed files with 34 additions and 30 deletions

View File

@ -4,6 +4,7 @@ import helpers, ssnet, ssh
from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
from helpers import *
def original_dst(sock):
try:
SO_ORIGINAL_DST = 80
@ -176,20 +177,7 @@ def _main(listener, fw, use_server, remotename, python, seed_hosts, auto_nets):
if rv:
raise Fatal('server died with error code %d' % rv)
r = []
w = []
x = []
handlers = filter(lambda s: s.ok, handlers)
for s in handlers:
s.pre_select(r,w,x)
debug2('Waiting: %d[%d,%d,%d]...\n'
% (len(handlers), len(r), len(w), len(x)))
(r,w,x) = select.select(r,w,x)
#log('r=%r w=%r x=%r\n' % (r,w,x))
ready = r+w+x
for s in handlers:
if list_contains_any(s.socks, ready):
s.callback()
ssnet.runonce(handlers, mux)
if use_server:
mux.callback()
mux.check_fullness()

View File

@ -162,21 +162,6 @@ def main():
if rpid:
raise Fatal('hostwatch exited unexpectedly: code 0x%04x\n' % rv)
r = []
w = []
x = []
handlers = filter(lambda s: s.ok, handlers)
for s in handlers:
s.pre_select(r,w,x)
debug2('Waiting: %d[%d,%d,%d] (fullness=%d/%d)...\n'
% (len(handlers), len(r), len(w), len(x),
mux.fullness, mux.too_full))
(r,w,x) = select.select(r,w,x)
#log('r=%r w=%r x=%r\n' % (r,w,x))
ready = r+w+x
for s in handlers:
#debug2('check: %r: %r\n' % (s, s.socks & ready))
if list_contains_any(s.socks, ready):
s.callback()
ssnet.runonce(handlers, mux)
mux.check_fullness()
mux.callback()

View File

@ -36,6 +36,17 @@ def _add(l, elem):
l.append(elem)
def _fds(l):
out = []
for i in l:
try:
out.append(i.fileno())
except AttributeError:
out.append(i)
out.sort()
return out
def _nb_clean(func, *args):
try:
return func(*args)
@ -43,6 +54,7 @@ def _nb_clean(func, *args):
if e.errno not in (errno.EWOULDBLOCK, errno.EAGAIN, errno.EPIPE):
raise
else:
debug3('%s: err was: %s\n' % (func.__name__, e))
return None
@ -429,3 +441,22 @@ def connect_dst(ip, port):
return SockWrapper(outsock, outsock,
connect_to = (ip,port),
peername = '%s:%d' % (ip,port))
def runonce(handlers, mux):
r = []
w = []
x = []
handlers = filter(lambda s: s.ok, handlers)
for s in handlers:
s.pre_select(r,w,x)
debug2('Waiting: %d r=%r w=%r x=%r (fullness=%d/%d)\n'
% (len(handlers), _fds(r), _fds(w), _fds(x),
mux.fullness, mux.too_full))
(r,w,x) = select.select(r,w,x)
debug2(' Ready: %d r=%r w=%r x=%r\n'
% (len(handlers), _fds(r), _fds(w), _fds(x)))
ready = r+w+x
for s in handlers:
if list_contains_any(s.socks, ready):
s.callback()