From 84376284db6199d4928797e5e8be4182449120f2 Mon Sep 17 00:00:00 2001
From: Avery Pennarun <apenwarr@gmail.com>
Date: Fri, 1 Oct 2010 17:36:09 -0700
Subject: [PATCH] Factor out common mainloop code between client and server.

Also improve the socket message output a bit.
---
 client.py | 16 ++--------------
 server.py | 17 +----------------
 ssnet.py  | 31 +++++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 30 deletions(-)

diff --git a/client.py b/client.py
index 0a4b2e5..4d657dd 100644
--- a/client.py
+++ b/client.py
@@ -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()
diff --git a/server.py b/server.py
index 31cd6c4..ae7c706 100644
--- a/server.py
+++ b/server.py
@@ -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()
diff --git a/ssnet.py b/ssnet.py
index 7aebe25..b8daf7d 100644
--- a/ssnet.py
+++ b/ssnet.py
@@ -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()