Add sock paramater to Handler callbacks

As Handler objects can have multiple sockets, we need to know which one
was involved in the incoming event.

3 files changed, 19 insertions, 19 deletions

diff --git a/sshuttle/client.py b/sshuttle/client.py
index 829838a..577d979 100644
--- a/sshuttle/client.py
+++ b/sshuttle/client.py
@@ -106,16 +106,18 @@ class MultiListener:
             self.v4.setsockopt(level, optname, value)
 
     def add_handler(self, handlers, callback, method, mux):
+        socks = []
         if self.v6:
-            handlers.append(
-                Handler(
-                    [self.v6],
-                    lambda: callback(self.v6, method, mux, handlers)))
+            socks.append(self.v6)
         if self.v4:
-            handlers.append(
-                Handler(
-                    [self.v4],
-                    lambda: callback(self.v4, method, mux, handlers)))
+            socks.append(self.v4)
+
+        handlers.append(
+            Handler(
+                socks,
+                lambda sock: callback(sock, method, mux, handlers)
+            )
+        )
 
     def listen(self, backlog):
         if self.v6:
@@ -480,7 +482,6 @@ def _main(tcp_listener, udp_listener, fw, ssh_cmd, remotename,
         ssnet.runonce(handlers, mux)
         if latency_control:
             mux.check_fullness()
-        mux.callback()
 
 
 def main(listenip_v6, listenip_v4,
diff --git a/sshuttle/server.py b/sshuttle/server.py
index 24b7d6b..2529765 100644
--- a/sshuttle/server.py
+++ b/sshuttle/server.py
@@ -156,9 +156,9 @@ class DnsProxy(Handler):
                 log('DNS send to %r: %s\n' % (self.peer, e))
                 return
 
-    def callback(self):
+    def callback(self, sock):
         try:
-            data = self.sock.recv(4096)
+            data = sock.recv(4096)
         except socket.error as e:
             if e.args[0] in ssnet.NET_ERRS:
                 # might have been spurious; try again.
@@ -195,9 +195,9 @@ class UdpProxy(Handler):
             log('UDP send to %r port %d: %s\n' % (dstip[0], dstip[1], e))
             return
 
-    def callback(self):
+    def callback(self, sock):
         try:
-            data, peer = self.sock.recvfrom(4096)
+            data, peer = sock.recvfrom(4096)
         except socket.error as e:
             log('UDP recv from %r port %d: %s\n' % (peer[0], peer[1], e))
             return
@@ -236,7 +236,7 @@ def main(latency_control):
     hw = Hostwatch()
     hw.leftover = ''
 
-    def hostwatch_ready():
+    def hostwatch_ready(sock):
         assert(hw.pid)
         content = hw.sock.recv(4096)
         if content:
@@ -314,7 +314,6 @@ def main(latency_control):
         ssnet.runonce(handlers, mux)
         if latency_control:
             mux.check_fullness()
-        mux.callback()
 
         if dnshandlers:
             now = time.time()
diff --git a/sshuttle/ssnet.py b/sshuttle/ssnet.py
index e20e1d5..218e7b5 100644
--- a/sshuttle/ssnet.py
+++ b/sshuttle/ssnet.py
@@ -260,7 +260,7 @@ class Handler:
         for i in self.socks:
             _add(r, i)
 
-    def callback(self):
+    def callback(self, sock):
         log('--no callback defined-- %r\n' % self)
         (r, w, x) = select.select(self.socks, [], [], 0)
         for s in r:
@@ -301,7 +301,7 @@ class Proxy(Handler):
         elif not self.wrap2.shut_read:
             _add(r, self.wrap2.rsock)
 
-    def callback(self):
+    def callback(self, sock):
         self.wrap1.try_connect()
         self.wrap2.try_connect()
         self.wrap1.fill()
@@ -467,7 +467,7 @@ class Mux(Handler):
         if self.outbuf:
             _add(w, self.wsock)
 
-    def callback(self):
+    def callback(self, sock):
         (r, w, x) = select.select([self.rsock], [self.wsock], [], 0)
         if self.rsock in r:
             self.handle()
@@ -572,7 +572,7 @@ def runonce(handlers, mux):
     for h in handlers:
         for s in h.socks:
             if s in ready:
-                h.callback()
+                h.callback(s)
                 did[s] = 1
     for s in ready:
         if s not in did: