From c8018f29c923df9ae82864df63f75d88a68ddb93 Mon Sep 17 00:00:00 2001 From: Karim Allah Ahmed Date: Tue, 22 May 2012 16:09:07 +0200 Subject: [PATCH 1/2] Adding support for proxying from a unix socket --- .gitignore | 2 ++ websocket.py | 54 +++++++++++++++++++++++++++++++++------------------- websockify | 40 ++++++++++++++++++++++---------------- 3 files changed, 60 insertions(+), 36 deletions(-) mode change 100755 => 100644 websockify diff --git a/.gitignore b/.gitignore index af62ef9..9c5203f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,5 @@ other/.lein-deps-sum other/classes other/lib +.project +.pydevproject diff --git a/websocket.py b/websocket.py index d3bb48c..6975248 100644 --- a/websocket.py +++ b/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,7 +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): """ 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. @@ -171,24 +174,30 @@ 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 not connect: flags = flags | socket.AI_PASSIVE - addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM, - socket.IPPROTO_TCP, flags) - if not addrs: - raise Exception("Could resolve host '%s'" % host) - addrs.sort(key=lambda x: x[0]) - if prefer_ipv6: - addrs.reverse() - sock = socket.socket(addrs[0][0], addrs[0][1]) - if connect: - sock.connect(addrs[0][4]) - else: - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind(addrs[0][4]) - sock.listen(100) + + if not unix_socket: + addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM, + socket.IPPROTO_TCP, flags) + if not addrs: + raise Exception("Could resolve host '%s'" % host) + addrs.sort(key=lambda x: x[0]) + if prefer_ipv6: + addrs.reverse() + sock = socket.socket(addrs[0][0], addrs[0][1]) + if connect: + sock.connect(addrs[0][4]) + else: + 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 @@ -569,10 +578,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 @@ -751,6 +760,11 @@ Sec-WebSocket-Accept: %s\r self.start_time = int(time.time()*1000) # handler process + + dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port) + dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port) + + try: try: self.client = self.do_handshake(startsock, address) @@ -848,7 +862,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) diff --git a/websockify b/websockify old mode 100755 new mode 100644 index 550dff7..ce12194 --- a/websockify +++ b/websockify @@ -88,14 +88,14 @@ Traffic Legend: # Need to call wrapped command after daemonization so we can # know when the wrapped command exits if self.wrap_cmd: - print(" - proxying from %s:%s to '%s' (port %s)\n" % ( - self.listen_host, self.listen_port, - " ".join(self.wrap_cmd), self.target_port)) + dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port) + print(" - proxying from %s:%s to %s\n" % ( + self.listen_host, self.listen_port, dst_string)) self.run_wrap_cmd() else: - print(" - proxying from %s:%s to %s:%s\n" % ( - self.listen_host, self.listen_port, - self.target_host, self.target_port)) + dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port) + print(" - proxying from %s:%s to %s\n" % ( + self.listen_host, self.listen_port, dst_string)) def poll(self): # If we are wrapping a command, check it's status @@ -137,12 +137,15 @@ Traffic Legend: """ Called after a new WebSocket connection has been established. """ - # Connect to the target - self.msg("connecting to: %s:%s" % ( - self.target_host, self.target_port)) + if self.unix_socket: + self.msg("connecting to unix socket : %s" % self.unix_socket) + else: + self.msg("connecting to: %s:%s" % ( + self.target_host, self.target_port)) + tsock = self.socket(self.target_host, self.target_port, - connect=True) + connect=True, unix_socket=self.unix_socket) if self.verbose and not self.daemon: print(self.traffic_legend) @@ -223,6 +226,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)") @@ -245,7 +250,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:] @@ -270,12 +275,15 @@ def websockify_init(): opts.target_host = None opts.target_port = None else: - if args[1].count(':') > 0: - opts.target_host, opts.target_port = args[1].rsplit(':', 1) + if hasattr(opts, 'unix'): + opts.target_host = opts.target_port = None else: - parser.error("Error parsing target") - try: opts.target_port = int(opts.target_port) - except: parser.error("Error parsing target port") + if args[1].count(':') > 0: + opts.target_host, opts.target_port = args[1].rsplit(':', 1) + 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__) From f3054df53a2fd538a4d2ee04ffed2dabbc904909 Mon Sep 17 00:00:00 2001 From: Karim Allah Ahmed Date: Tue, 22 May 2012 16:49:00 +0200 Subject: [PATCH 2/2] Adding TCP_NODELAY to the source sockets --- websocket.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/websocket.py b/websocket.py index 6975248..9af0424 100644 --- a/websocket.py +++ b/websocket.py @@ -759,12 +759,7 @@ Sec-WebSocket-Accept: %s\r self.rec = None self.start_time = int(time.time()*1000) - # handler process - - dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port) - dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port) - - + # handler process try: try: self.client = self.do_handshake(startsock, address)