Add --wrap option to force channel number wrapping at a lower number.

This makes it easier to actually test what happens when channel numbers wrap
around.  The good news: it works.

However, I did find a bug where sshuttle would die if we completely ran out
of available channel numbers because so many of them were open.  This would
never realistically happen at the default of 65535 channels (we'd run out of
file descriptors first), but it's still a bug, so let's handle it by just
dropping the connection when it happens.

3 files changed, 11 insertions, 1 deletions

diff --git a/client.py b/client.py
index 1ade5d9..fa93c26 100644
--- a/client.py
+++ b/client.py
@@ -273,6 +273,10 @@ def _main(listener, fw, ssh_cmd, remotename, python, latency_control,
             sock.close()
             return
         chan = mux.next_channel()
+        if not chan:
+            log('warning: too many open channels.  Discarded connection.\n')
+            sock.close()
+            return
         mux.send(chan, ssnet.CMD_CONNECT, '%s,%s' % dstip)
         outwrap = MuxWrapper(mux, chan)
         handlers.append(Proxy(SockWrapper(sock, sock), outwrap))
diff --git a/main.py b/main.py
index e76e596..3fe1e53 100755
--- a/main.py
+++ b/main.py
@@ -62,6 +62,7 @@ v,verbose  increase debug message verbosity
 e,ssh-cmd= the command to use to connect to the remote [ssh]
 seed-hosts= with -H, use these hostnames for initial scan (comma-separated)
 no-latency-control  sacrifice latency to improve bandwidth benchmarks
+wrap=      restart counting channel numbers after this number (for testing)
 D,daemon   run in the background as a daemon
 syslog     send log messages to syslog (default if you use --daemon)
 pidfile=   pidfile name (only if using --daemon) [./sshuttle.pid]
@@ -74,6 +75,9 @@ o = options.Options(optspec)
 
 if opt.daemon:
     opt.syslog = 1
+if opt.wrap:
+    import ssnet
+    ssnet.MAX_CHANNEL = int(opt.wrap)
 helpers.verbose = opt.verbose
 
 try:
diff --git a/ssnet.py b/ssnet.py
index 554d870..2145431 100644
--- a/ssnet.py
+++ b/ssnet.py
@@ -1,6 +1,8 @@
 import struct, socket, errno, select
 if not globals().get('skip_imports'):
     from helpers import *
+
+MAX_CHANNEL = 65535
     
 # these don't exist in the socket module in python 2.3!
 SHUT_RD = 0
@@ -300,7 +302,7 @@ class Mux(Handler):
         # channel 0 is special, so we never allocate it
         for timeout in xrange(1024):
             self.chani += 1
-            if self.chani > 65535:
+            if self.chani > MAX_CHANNEL:
                 self.chani = 1
             if not self.channels.get(self.chani):
                 return self.chani