debug + work for the next connections by not closing the accepter port on None handle_input

2 files changed, 29 insertions, 16 deletions

diff --git a/bandwidthmon.py b/bandwidthmon.py
index 4e88c16..c99aaac 100755
--- a/bandwidthmon.py
+++ b/bandwidthmon.py
@@ -6,6 +6,11 @@ import glib
 import gtk
 
 def main():
+    debug = True
+    def dprint(x):
+        if not debug:
+            return
+        print x
     p = argparse.ArgumentParser(description="proxy multiple socket connections")
     p.add_argument('--listen-port', required=True, type=int)
     p.add_argument('--remote-port', required=True, type=int)
@@ -14,10 +19,12 @@ def main():
     args = p.parse_args(sys.argv[1:])
     iterate_packets, handle_input, select_based_iterator, get_fds = proxy.make_proxy(
         args.listen_port, (args.remote_host, args.remote_port),
-        args.listen_host, debug=True)
+        args.listen_host, debug=debug)
+    assert(len(get_fds()) == 1) # only the accepting socket
+    accepter = get_fds()[0]
     added_fds = set()
     def on_new_fd(fd):
-        print "on_new_fd %s" % fd.fileno()
+        dprint("on_new_fd %s" % fd.fileno())
         # Don't add glib.IO_OUT unless you mean it.
         # corollary: when I want to really implement this (to implement non
         # blocking writes) I'll have to io_add_watch(OUT) and return False and
@@ -30,26 +37,26 @@ def main():
         try:
             fd.fileno()
         except socket.error:
-            print "removing socket %s" % fd
+            dprint("removing socket %s" % fd)
             added_fds.remove(fd)
             return False
-        print "called back %s, condition %s" % (fd.fileno(), condition)
+        dprint("called back %s, condition %s" % (fd.fileno(), condition))
         if condition != glib.IO_IN:
-            print "not reading from %d" % fd.fileno()
+            dprint("not reading from %d" % fd.fileno())
             return True
         result = handle_input(fd)
-        print "%s: result %s" % (fd.fileno(), repr(result))
+        dprint("%s: result %s" % (fd.fileno(), repr(result)))
         fds, new_fds = update_fds()
         if not result:
-            # closed port
-            return False
+            # accepter port returns nothing, don't close it.
+            return fd == accepter
         (src, dst, data, completer) = result
         completer()
         print repr((src, dst, data))
         return True
     def update_fds():
         fds = get_fds()
-        print "update_fds: %s" % (','.join(str(f.fileno()) for f in fds))
+        dprint("update_fds: %s" % (','.join(str(f.fileno()) for f in fds)))
         new_fds = []
         for fd in set(fds) - added_fds:
             on_new_fd(fd)
diff --git a/proxy.py b/proxy.py
index 1d7327c..1e53a10 100755
--- a/proxy.py
+++ b/proxy.py
@@ -61,7 +61,7 @@ class Proxy(object):
 
     def __init__(self, local_port, remote_addr, host='127.0.0.1'):
         self._drop_next = False
-        handle_input, select_based_iterator, get_fds = make_proxy(
+        iterate_packets, handle_input, select_based_iterator, get_fds = make_proxy(
             local_port, remote_addr, host, check_drop_next=self.check_drop_next)
         self.select_based_iterator = select_based_iterator()
 
@@ -87,7 +87,7 @@ def make_proxy(local_port, remote_addr, host = '127.0.0.1',
     def get_fds():
         fds = [accepter] + open_socks.allkeys()
         if debug:
-            print repr(fds)
+            print "make_proxy:", repr(fds)
         return fds
     def iterate_packets():
         """
@@ -138,7 +138,13 @@ def make_proxy(local_port, remote_addr, host = '127.0.0.1',
             if dont_recv:
                 return None
             src, dst = src_dst
-            data = s.recv(MAX_PACKET_SIZE)
+            try:
+                if debug:
+                    print "recv from %d" % s.fileno()
+                data = s.recv(MAX_PACKET_SIZE)
+            except socket.error:
+                print "handle_input: bad event loop - socket invalid %s" % s.fileno()
+                return None
             if len(data) == 0:
                 other.close()
                 del open_socks[s]
@@ -167,7 +173,7 @@ def make_proxy(local_port, remote_addr, host = '127.0.0.1',
             results = packet_iter.send(rds)
             for r in results:
                 yield r
-    return handle_input, select_based_iterator, get_fds
+    return iterate_packets, handle_input, select_based_iterator, get_fds
 
 def proxy(local_port, remote_addr):
     return Proxy(local_port, remote_addr)
@@ -190,12 +196,12 @@ def tests():
     import argparse
     import sys
     p = argparse.ArgumentParser(description="proxy multiple socket connections")
-    p.add_argument('--listen-port', default=11000)
+    p.add_argument('--listen-port', default=11000, type=int)
     p.add_argument('--listen-host', default='127.0.0.1')
-    p.add_argument('--remote-port', default=12000)
+    p.add_argument('--remote-port', default=12000, type=int)
     p.add_argument('--remote-host', default='127.0.0.1')
     args = p.parse_args(sys.argv[1:])
-    handle_input, select_based_iterator, get_fds = make_proxy(
+    iterate_packets, handle_input, select_based_iterator, get_fds = make_proxy(
         args.listen_port, (args.remote_host, args.remote_port),
         args.listen_host, debug=True)
     for x in select_based_iterator():