Merge branch 'closing'

* closing:
  Correctly close server connection when client disconnects.
  "Too many open files" shouldn't be a fatal condition.

2 files changed, 29 insertions, 4 deletions

diff --git a/client.py b/client.py
index a1ceaae..595c42a 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 *
 
+_extra_fd = os.open('/dev/null', os.O_RDONLY)
 
 def _islocal(ip):
     sock = socket.socket()
@@ -168,7 +169,22 @@ def _main(listener, fw, ssh_cmd, remotename, python, seed_hosts, auto_nets):
     mux.got_host_list = onhostlist
 
     def onaccept():
-        sock,srcip = listener.accept()
+        global _extra_fd
+        try:
+            sock,srcip = listener.accept()
+        except socket.error, e:
+            if e.args[0] in [errno.EMFILE, errno.ENFILE]:
+                debug1('Rejected incoming connection: too many open files!\n')
+                # free up an fd so we can eat the connection
+                os.close(_extra_fd)
+                try:
+                    sock,srcip = listener.accept()
+                    sock.close()
+                finally:
+                    _extra_fd = os.open('/dev/null', os.O_RDONLY)
+                return
+            else:
+                raise
         dstip = original_dst(sock)
         debug1('Accept: %s:%r -> %s:%r.\n' % (srcip[0],srcip[1],
                                               dstip[0],dstip[1]))
diff --git a/ssnet.py b/ssnet.py
index c06498f..c325d76 100644
--- a/ssnet.py
+++ b/ssnet.py
@@ -15,7 +15,7 @@ CMD_EXIT = 0x4200
 CMD_PING = 0x4201
 CMD_PONG = 0x4202
 CMD_CONNECT = 0x4203
-# CMD_CLOSE = 0x4204   # never used - removed
+CMD_STOP_SENDING = 0x4204
 CMD_EOF = 0x4205
 CMD_DATA = 0x4206
 CMD_ROUTES = 0x4207
@@ -27,6 +27,7 @@ cmd_to_name = {
     CMD_PING: 'PING',
     CMD_PONG: 'PONG',
     CMD_CONNECT: 'CONNECT',
+    CMD_STOP_SENDING: 'STOP_SENDING',
     CMD_EOF: 'EOF',
     CMD_DATA: 'DATA',
     CMD_ROUTES: 'ROUTES',
@@ -106,6 +107,8 @@ class SockWrapper:
     def seterr(self, e):
         if not self.exc:
             self.exc = e
+        self.nowrite()
+        self.noread()
 
     def try_connect(self):
         if self.connect_to and self.shut_write:
@@ -162,8 +165,6 @@ class SockWrapper:
         except OSError, e:
             # unexpected error... stream is dead
             self.seterr(e)
-            self.nowrite()
-            self.noread()
             return 0
         
     def write(self, buf):
@@ -231,6 +232,11 @@ class Proxy(Handler):
         self.wrap2 = wrap2
 
     def pre_select(self, r, w, x):
+        if self.wrap1.shut_read: self.wrap2.nowrite()
+        if self.wrap1.shut_write: self.wrap2.noread()
+        if self.wrap2.shut_read: self.wrap1.nowrite()
+        if self.wrap2.shut_write: self.wrap1.noread()
+        
         if self.wrap1.connect_to:
             _add(w, self.wrap1.rsock)
         elif self.wrap1.buf:
@@ -430,6 +436,7 @@ class MuxWrapper(SockWrapper):
     def noread(self):
         if not self.shut_read:
             self.shut_read = True
+            self.mux.send(self.channel, CMD_STOP_SENDING, '')
             self.maybe_close()
 
     def nowrite(self):
@@ -464,6 +471,8 @@ class MuxWrapper(SockWrapper):
     def got_packet(self, cmd, data):
         if cmd == CMD_EOF:
             self.noread()
+        elif cmd == CMD_STOP_SENDING:
+            self.nowrite()
         elif cmd == CMD_DATA:
             self.buf.append(data)
         else: