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

1 files changed, 13 insertions, 7 deletions

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():