Merge pull request #46 from cloud9ers/master

Adding support for unix sockets as target socket.

.gitignore

@@ -4,3 +4,5 @@
 other/.lein-deps-sum
 other/classes
 other/lib
+.project
+.pydevproject

websocket.py

@@ -75,6 +75,7 @@ class WebSocketServer(object):
 
     buffer_size = 65536
 
+
     server_handshake_hixie = """HTTP/1.1 101 Web Socket Protocol Handshake\r
 Upgrade: WebSocket\r
 Connection: Upgrade\r
@@ -103,7 +104,7 @@ Sec-WebSocket-Accept: %s\r
     def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
             verbose=False, cert='', key='', ssl_only=None,
             daemon=False, record='', web='',
-            run_once=False, timeout=0):
+            run_once=False, timeout=0, unix=None):
 
         # settings
         self.verbose        = verbose
@@ -113,6 +114,8 @@ Sec-WebSocket-Accept: %s\r
         self.daemon         = daemon
         self.run_once       = run_once
         self.timeout        = timeout
+        
+        self.unix_socket    = unix
 
         self.launch_time    = time.time()
         self.ws_connection  = False
@@ -163,8 +166,7 @@ Sec-WebSocket-Accept: %s\r
     #
 
     @staticmethod
-    def socket(host, port=None, connect=False, prefer_ipv6=False,
+    def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None, use_ssl=False):
-               use_ssl=False):
         """ Resolve a host (and optional port) to an IPv4 or IPv6
         address. Create a socket. Bind to it if listen is set,
         otherwise connect to it. Return the socket.
@@ -172,7 +174,7 @@ Sec-WebSocket-Accept: %s\r
         flags = 0
         if host == '':
             host = None
-        if connect and not port:
+        if connect and not (port or unix_socket):
             raise Exception("Connect mode requires a port")
         if use_ssl and not ssl:
             raise Exception("SSL socket requested but Python SSL module not loaded.");
@@ -180,22 +182,28 @@ Sec-WebSocket-Accept: %s\r
             raise Exception("SSL only supported in connect mode (for now)")
         if not connect:
             flags = flags | socket.AI_PASSIVE
-        addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
+            
-                socket.IPPROTO_TCP, flags)
+        if not unix_socket:
-        if not addrs:
+            addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
-            raise Exception("Could resolve host '%s'" % host)
+                    socket.IPPROTO_TCP, flags)
-        addrs.sort(key=lambda x: x[0])
+            if not addrs:
-        if prefer_ipv6:
+                raise Exception("Could not resolve host '%s'" % host)
-            addrs.reverse()
+            addrs.sort(key=lambda x: x[0])
-        sock = socket.socket(addrs[0][0], addrs[0][1])
+            if prefer_ipv6:
-        if connect:
+                addrs.reverse()
-            sock.connect(addrs[0][4])
+            sock = socket.socket(addrs[0][0], addrs[0][1])
-            if use_ssl:
+            if connect:
-                sock = ssl.wrap_socket(sock)
+                sock.connect(addrs[0][4])
-        else:
+                if use_ssl:
-            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                    sock = ssl.wrap_socket(sock)
-            sock.bind(addrs[0][4])
+            else:
-            sock.listen(100)
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                sock.bind(addrs[0][4])
+                sock.listen(100)
+        else:    
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            sock.connect(unix_socket)
+
         return sock
 
     @staticmethod
@@ -576,10 +584,10 @@ Sec-WebSocket-Accept: %s\r
         - Send a WebSockets handshake server response.
         - Return the socket for this WebSocket client.
         """
-
         stype = ""
-
         ready = select.select([sock], [], [], 3)[0]
+
+        
         if not ready:
             raise self.EClose("ignoring socket not ready")
         # Peek, but do not read the data so that we have a opportunity
@@ -757,7 +765,7 @@ Sec-WebSocket-Accept: %s\r
         self.rec        = None
         self.start_time = int(time.time()*1000)
 
-        # handler process
+        # handler process        
         try:
             try:
                 self.client = self.do_handshake(startsock, address)
@@ -855,7 +863,7 @@ Sec-WebSocket-Accept: %s\r
                             continue
                         else:
                             raise
-
+                    
                     if self.run_once:
                         # Run in same process if run_once
                         self.top_new_client(startsock, address)

websockify

@@ -88,13 +88,15 @@ Traffic Legend:
         """
         # Need to call wrapped command after daemonization so we can
         # know when the wrapped command exits
-        msg = "  - proxying from %s:%s" % (
-                self.listen_host, self.listen_port)
         if self.wrap_cmd:
-            msg += " to '%s' - port %s" % (
+            dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
-                    " ".join(self.wrap_cmd, self.target_port))
+            msg = "  - proxying from %s:%s to %s\n" % (
+                    self.listen_host, self.listen_port, dst_string)
+            self.run_wrap_cmd()
         else:
-            msg += " to %s:%s" % (self.target_host, self.listen_port)
+            dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port)
+            msg = "  - proxying from %s:%s to %s\n" % (
+                    self.listen_host, self.listen_port, dst_string)
 
         if self.ssl_target:
             msg += " (using SSL)"
@@ -144,11 +146,16 @@ Traffic Legend:
         """
         Called after a new WebSocket connection has been established.
         """
-
         # Connect to the target
+        if self.unix_socket:
+            msg = "connecting to unix socket : %s" % self.unix_socket
+        else:
+            msg = "connecting to: %s:%s" % (
+                                    self.target_host, self.target_port)
+        
+        tsock = self.socket(self.target_host, self.target_port,
+                            connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_socket)
          
-        msg = "connecting to: %s:%s" % (
-                self.target_host, self.target_port)
         if self.ssl_target:
             msg += " (using SSL)"
         self.msg(msg)
@@ -235,6 +242,8 @@ def websockify_init():
             help="verbose messages and per frame traffic")
     parser.add_option("--record",
             help="record sessions to FILE.[session_number]", metavar="FILE")
+    parser.add_option("--unix",
+            help="unix socket to proxy network from", metavar="FILE")
     parser.add_option("--daemon", "-D",
             dest="daemon", action="store_true",
             help="become a daemon (background process)")
@@ -259,7 +268,7 @@ def websockify_init():
     (opts, args) = parser.parse_args()
 
     # Sanity checks
-    if len(args) < 2:
+    if len(args) < 1:
         parser.error("Too few arguments")
     if sys.argv.count('--'):
         opts.wrap_cmd = args[1:]
@@ -287,12 +296,15 @@ def websockify_init():
         opts.target_host = None
         opts.target_port = None
     else:
-        if args[1].count(':') > 0:
+        if hasattr(opts, 'unix'):
-            opts.target_host, opts.target_port = args[1].rsplit(':', 1)
+            opts.target_host = opts.target_port = None
         else:
-            parser.error("Error parsing target")
+            if args[1].count(':') > 0:
-        try:    opts.target_port = int(opts.target_port)
+                opts.target_host, opts.target_port = args[1].rsplit(':', 1)
-        except: parser.error("Error parsing target port")
+            else:
+                parser.error("Error parsing target")
+            try:    opts.target_port = int(opts.target_port)
+            except: parser.error("Error parsing target port")
 
     # Create and start the WebSockets proxy
     server = WebSocketProxy(**opts.__dict__)