Merge db93395061 into 34a1b68d79

tests/echo.py

@@ -12,15 +12,15 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
 
 import os, sys, select, optparse
 sys.path.insert(0,os.path.join(os.path.dirname(__file__), ".."))
-from websockify.websocket import WebSocketServer
+from websockify.websocket import WebSocketServer, WebSocketRequestHandler
 
-class WebSocketEcho(WebSocketServer):
+class WebSocketEcho(WebSocketRequestHandler):
     """
     WebSockets server that echos back whatever is received from the
     client.  """
     buffer_size = 8096
 
-    def new_client(self):
+    def new_websocket_client(self):
         """
         Echo back whatever is received.
         """
@@ -28,28 +28,27 @@ class WebSocketEcho(WebSocketServer):
         cqueue = []
         c_pend = 0
         cpartial = ""
-        rlist = [self.client]
+        rlist = [self.request]
 
         while True:
             wlist = []
 
-            if cqueue or c_pend: wlist.append(self.client)
+            if cqueue or c_pend: wlist.append(self.request)
             ins, outs, excepts = select.select(rlist, wlist, [], 1)
             if excepts: raise Exception("Socket exception")
 
-            if self.client in outs:
+            if self.request in outs:
                 # Send queued target data to the client
                 c_pend = self.send_frames(cqueue)
                 cqueue = []
 
-            if self.client in ins:
+            if self.request in ins:
                 # Receive client data, decode it, and send it back
                 frames, closed = self.recv_frames()
                 cqueue.extend(frames)
 
                 if closed:
                     self.send_close()
-                    raise self.EClose(closed)
 
 if __name__ == '__main__':
     parser = optparse.OptionParser(usage="%prog [options] listen_port")
@@ -70,6 +69,6 @@ if __name__ == '__main__':
         parser.error("Invalid arguments")
 
     opts.web = "."
-    server = WebSocketEcho(**opts.__dict__)
+    server = WebSocketServer(WebSocketEcho, **opts.__dict__)
     server.start_server()
 

tests/load.py

@@ -8,33 +8,35 @@ given a sequence number. Any errors are reported and counted.
 
 import sys, os, select, random, time, optparse
 sys.path.insert(0,os.path.join(os.path.dirname(__file__), ".."))
-from websockify.websocket import WebSocketServer
+from websockify.websocket import WebSocketServer, WebSocketRequestHandler
 
-class WebSocketLoad(WebSocketServer):
+class WebSocketLoadServer(WebSocketServer):
 
-    buffer_size = 65536
-
-    max_packet_size = 10000
     recv_cnt = 0
     send_cnt = 0
 
     def __init__(self, *args, **kwargs):
-        self.errors = 0
         self.delay = kwargs.pop('delay')
 
+        WebSocketServer.__init__(self, *args, **kwargs)
+
+
+class WebSocketLoad(WebSocketRequestHandler):
+
+    max_packet_size = 10000
+
+    def new_websocket_client(self):
         print "Prepopulating random array"
         self.rand_array = []
         for i in range(0, self.max_packet_size):
             self.rand_array.append(random.randint(0, 9))
 
-        WebSocketServer.__init__(self, *args, **kwargs)
+        self.errors = 0
-
-    def new_client(self):
         self.send_cnt = 0
         self.recv_cnt = 0
 
         try:
-            self.responder(self.client)
+            self.responder(self.request)
         except:
             print "accumulated errors:", self.errors
             self.errors = 0
@@ -61,14 +63,13 @@ class WebSocketLoad(WebSocketServer):
 
                 if closed:
                     self.send_close()
-                    raise self.EClose(closed)
 
             now = time.time() * 1000
             if client in outs:
                 if c_pend:
                     last_send = now
                     c_pend = self.send_frames()
-                elif now > (last_send + self.delay):
+                elif now > (last_send + self.server.delay):
                     last_send = now
                     c_pend = self.send_frames([self.generate()])
 
@@ -162,6 +163,6 @@ if __name__ == '__main__':
         parser.error("Invalid arguments")
 
     opts.web = "."
-    server = WebSocketLoad(**opts.__dict__)
+    server = WebSocketLoadServer(WebSocketLoad, **opts.__dict__)
     server.start_server()
 

websockify/websocketproxy.py

@@ -12,6 +12,10 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
 '''
 
 import signal, socket, optparse, time, os, sys, subprocess, logging
+try:    from socketserver import ForkingMixIn
+except: from SocketServer import ForkingMixIn
+try:    from http.server import HTTPServer
+except: from BaseHTTPServer import HTTPServer
 from select import select
 import websocket
 try:
@@ -20,15 +24,7 @@ except:
     from cgi import parse_qs
     from urlparse import urlparse
 
-class WebSocketProxy(websocket.WebSocketServer):
+class ProxyRequestHandler(websocket.WebSocketRequestHandler):
-    """
-    Proxy traffic to and from a WebSockets client to a normal TCP
-    socket server target. All traffic to/from the client is base64
-    encoded/decoded to allow binary data to be sent/received to/from
-    the target.
-    """
-
-    buffer_size = 65536
 
     traffic_legend = """
 Traffic Legend:
@@ -42,7 +38,156 @@ Traffic Legend:
     <. - Client send partial
 """
 
-    def __init__(self, *args, **kwargs):
+    def new_websocket_client(self):
+        """
+        Called after a new WebSocket connection has been established.
+        """
+        # Checks if we receive a token, and look
+        # for a valid target for it then
+        if self.server.target_cfg:
+            (self.server.target_host, self.server.target_port) = self.get_target(self.server.target_cfg, self.path)
+
+        # Connect to the target
+        if self.server.wrap_cmd:
+            msg = "connecting to command: '%s' (port %s)" % (" ".join(self.server.wrap_cmd), self.server.target_port)
+        elif self.server.unix_target:
+            msg = "connecting to unix socket: %s" % self.server.unix_target
+        else:
+            msg = "connecting to: %s:%s" % (
+                                    self.server.target_host, self.server.target_port)
+
+        if self.server.ssl_target:
+            msg += " (using SSL)"
+        self.log_message(msg)
+
+        tsock = websocket.WebSocketServer.socket(self.server.target_host,
+                                                 self.server.target_port,
+                connect=True, use_ssl=self.server.ssl_target, unix_socket=self.server.unix_target)
+
+        self.print_traffic(self.traffic_legend)
+
+        # Start proxying
+        try:
+            self.do_proxy(tsock)
+        except:
+            if tsock:
+                tsock.shutdown(socket.SHUT_RDWR)
+                tsock.close()
+                if self.verbose: 
+                    self.log_message("%s:%s: Closed target" %(
+                            self.server.target_host, self.server.target_port))
+            raise
+
+    def get_target(self, target_cfg, path):
+        """
+        Parses the path, extracts a token, and looks for a valid
+        target for that token in the configuration file(s). Sets
+        target_host and target_port if successful
+        """
+        # The files in targets contain the lines
+        # in the form of token: host:port
+
+        # Extract the token parameter from url
+        args = parse_qs(urlparse(path)[4]) # 4 is the query from url
+
+        if not args.has_key('token') or not len(args['token']):
+            raise self.EClose("Token not present")
+
+        token = args['token'][0].rstrip('\n')
+
+        # target_cfg can be a single config file or directory of
+        # config files
+        if os.path.isdir(target_cfg):
+            cfg_files = [os.path.join(target_cfg, f)
+                         for f in os.listdir(target_cfg)]
+        else:
+            cfg_files = [target_cfg]
+
+        targets = {}
+        for f in cfg_files:
+            for line in [l.strip() for l in file(f).readlines()]:
+                if line and not line.startswith('#'):
+                    ttoken, target = line.split(': ')
+                    targets[ttoken] = target.strip()
+
+        self.vmsg("Target config: %s" % repr(targets))
+
+        if targets.has_key(token):
+            return targets[token].split(':')
+        else:
+            raise self.EClose("Token '%s' not found" % token)
+
+    def do_proxy(self, target):
+        """
+        Proxy client WebSocket to normal target socket.
+        """
+        cqueue = []
+        c_pend = 0
+        tqueue = []
+        rlist = [self.request, target]
+
+        while True:
+            wlist = []
+
+            if tqueue: wlist.append(target)
+            if cqueue or c_pend: wlist.append(self.request)
+            ins, outs, excepts = select(rlist, wlist, [], 1)
+            if excepts: raise Exception("Socket exception")
+
+            if self.request in outs:
+                # Send queued target data to the client
+                c_pend = self.send_frames(cqueue)
+
+                cqueue = []
+
+            if self.request in ins:
+                # Receive client data, decode it, and queue for target
+                bufs, closed = self.recv_frames()
+                tqueue.extend(bufs)
+
+                if closed:
+                    # TODO: What about blocking on client socket?
+                    if self.verbose: 
+                        self.log_message("%s:%s: Client closed connection" %(
+                                self.server.target_host, self.server.target_port))
+                    raise self.CClose(closed['code'], closed['reason'])
+
+
+            if target in outs:
+                # Send queued client data to the target
+                dat = tqueue.pop(0)
+                sent = target.send(dat)
+                if sent == len(dat):
+                    self.print_traffic(">")
+                else:
+                    # requeue the remaining data
+                    tqueue.insert(0, dat[sent:])
+                    self.print_traffic(".>")
+
+
+            if target in ins:
+                # Receive target data, encode it and queue for client
+                buf = target.recv(self.buffer_size)
+                if len(buf) == 0:
+                    if self.verbose:
+                        self.log_message("%s:%s: Target closed connection" %(
+                                self.server.target_host, self.server.target_port))
+                    raise self.CClose(1000, "Target closed")
+
+                cqueue.append(buf)
+                self.print_traffic("{")
+
+class WebSocketProxy(websocket.WebSocketServer):
+    """
+    Proxy traffic to and from a WebSockets client to a normal TCP
+    socket server target. All traffic to/from the client is base64
+    encoded/decoded to allow binary data to be sent/received to/from
+    the target.
+    """
+
+    buffer_size = 65536
+
+    def __init__(self, RequestHandlerClass=ProxyRequestHandler, *args, **kwargs):
         # Save off proxy specific options
         self.target_host    = kwargs.pop('target_host', None)
         self.target_port    = kwargs.pop('target_port', None)
@@ -83,7 +228,7 @@ Traffic Legend:
                 "REBIND_OLD_PORT": str(kwargs['listen_port']),
                 "REBIND_NEW_PORT": str(self.target_port)})
 
-        websocket.WebSocketServer.__init__(self, *args, **kwargs)
+        websocket.WebSocketServer.__init__(self, RequestHandlerClass, *args, **kwargs)
 
     def run_wrap_cmd(self):
         self.msg("Starting '%s'", " ".join(self.wrap_cmd))
@@ -147,152 +292,6 @@ Traffic Legend:
                 else:
                     self.run_wrap_cmd()
 
-    #
-    # Routines above this point are run in the master listener
-    # process.
-    #
-
-    #
-    # Routines below this point are connection handler routines and
-    # will be run in a separate forked process for each connection.
-    #
-
-    def new_client(self):
-        """
-        Called after a new WebSocket connection has been established.
-        """
-        # Checks if we receive a token, and look
-        # for a valid target for it then
-        if self.target_cfg:
-            (self.target_host, self.target_port) = self.get_target(self.target_cfg, self.path)
-
-        # Connect to the target
-        if self.wrap_cmd:
-            msg = "connecting to command: '%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
-        elif self.unix_target:
-            msg = "connecting to unix socket: %s" % self.unix_target
-        else:
-            msg = "connecting to: %s:%s" % (
-                                    self.target_host, self.target_port)
-
-        if self.ssl_target:
-            msg += " (using SSL)"
-        self.msg(msg)
-
-        tsock = self.socket(self.target_host, self.target_port,
-                connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_target)
-
-        self.print_traffic(self.traffic_legend)
-
-        # Start proxying
-        try:
-            self.do_proxy(tsock)
-        except:
-            if tsock:
-                tsock.shutdown(socket.SHUT_RDWR)
-                tsock.close()
-                self.vmsg("%s:%s: Closed target" %(
-                    self.target_host, self.target_port))
-            raise
-
-    def get_target(self, target_cfg, path):
-        """
-        Parses the path, extracts a token, and looks for a valid
-        target for that token in the configuration file(s). Sets
-        target_host and target_port if successful
-        """
-        # The files in targets contain the lines
-        # in the form of token: host:port
-
-        # Extract the token parameter from url
-        args = parse_qs(urlparse(path)[4]) # 4 is the query from url
-
-        if not args.has_key('token') or not len(args['token']):
-            raise self.EClose("Token not present")
-
-        token = args['token'][0].rstrip('\n')
-
-        # target_cfg can be a single config file or directory of
-        # config files
-        if os.path.isdir(target_cfg):
-            cfg_files = [os.path.join(target_cfg, f)
-                         for f in os.listdir(target_cfg)]
-        else:
-            cfg_files = [target_cfg]
-
-        targets = {}
-        for f in cfg_files:
-            for line in [l.strip() for l in file(f).readlines()]:
-                if line and not line.startswith('#'):
-                    ttoken, target = line.split(': ')
-                    targets[ttoken] = target.strip()
-
-        self.vmsg("Target config: %s" % repr(targets))
-
-        if targets.has_key(token):
-            return targets[token].split(':')
-        else:
-            raise self.EClose("Token '%s' not found" % token)
-
-    def do_proxy(self, target):
-        """
-        Proxy client WebSocket to normal target socket.
-        """
-        cqueue = []
-        c_pend = 0
-        tqueue = []
-        rlist = [self.client, target]
-
-        while True:
-            wlist = []
-
-            if tqueue: wlist.append(target)
-            if cqueue or c_pend: wlist.append(self.client)
-            ins, outs, excepts = select(rlist, wlist, [], 1)
-            if excepts: raise Exception("Socket exception")
-
-            if self.client in outs:
-                # Send queued target data to the client
-                c_pend = self.send_frames(cqueue)
-
-                cqueue = []
-
-            if self.client in ins:
-                # Receive client data, decode it, and queue for target
-                bufs, closed = self.recv_frames()
-                tqueue.extend(bufs)
-
-                if closed:
-                    # TODO: What about blocking on client socket?
-                    self.vmsg("%s:%s: Client closed connection" %(
-                        self.target_host, self.target_port))
-                    raise self.CClose(closed['code'], closed['reason'])
-
-
-            if target in outs:
-                # Send queued client data to the target
-                dat = tqueue.pop(0)
-                sent = target.send(dat)
-                if sent == len(dat):
-                    self.print_traffic(">")
-                else:
-                    # requeue the remaining data
-                    tqueue.insert(0, dat[sent:])
-                    self.print_traffic(".>")
-
-
-            if target in ins:
-                # Receive target data, encode it and queue for client
-                buf = target.recv(self.buffer_size)
-                if len(buf) == 0:
-                    self.vmsg("%s:%s: Target closed connection" %(
-                        self.target_host, self.target_port))
-                    raise self.CClose(1000, "Target closed")
-
-                cqueue.append(buf)
-                self.print_traffic("{")
-
-
 
 def _subprocess_setup():
     # Python installs a SIGPIPE handler by default. This is usually not what
@@ -358,6 +357,8 @@ def websockify_init():
             help="Configuration file containing valid targets "
             "in the form 'token: host:port' or, alternatively, a "
             "directory containing configuration files of this form")
+    parser.add_option("--libserver", action="store_true",
+            help="use Python library SocketServer engine")
     (opts, args) = parser.parse_args()
 
     if opts.verbose:
@@ -406,8 +407,65 @@ def websockify_init():
         opts.target_cfg = os.path.abspath(opts.target_cfg)
 
     # Create and start the WebSockets proxy
-    server = WebSocketProxy(**opts.__dict__)
+    libserver = opts.libserver
-    server.start_server()
+    del opts.libserver
+    if libserver:
+        # Use standard Python SocketServer framework
+        server = LibProxyServer(**opts.__dict__)
+        server.serve_forever()
+    else:
+        # Use internal service framework
+        server = WebSocketProxy(**opts.__dict__)
+        server.start_server()
+
+
+class LibProxyServer(ForkingMixIn, HTTPServer):
+    """
+    Just like WebSocketProxy, but uses standard Python SocketServer
+    framework.
+    """
+
+    def __init__(self, RequestHandlerClass=ProxyRequestHandler, **kwargs):
+        # Save off proxy specific options
+        self.target_host    = kwargs.pop('target_host', None)
+        self.target_port    = kwargs.pop('target_port', None)
+        self.wrap_cmd       = kwargs.pop('wrap_cmd', None)
+        self.wrap_mode      = kwargs.pop('wrap_mode', None)
+        self.unix_target    = kwargs.pop('unix_target', None)
+        self.ssl_target     = kwargs.pop('ssl_target', None)
+        self.target_cfg     = kwargs.pop('target_cfg', None)
+        self.daemon = False
+        self.target_cfg = None
+
+        # Server configuration
+        listen_host    = kwargs.pop('listen_host', '')
+        listen_port    = kwargs.pop('listen_port', None)
+        web            = kwargs.pop('web', '')
+
+        # Configuration affecting base request handler
+        self.only_upgrade   = not web
+        self.verbose   = kwargs.pop('verbose', False)
+        record = kwargs.pop('record', '')
+        if record:
+            self.record = os.path.abspath(record)
+        self.run_once  = kwargs.pop('run_once', False)
+        self.handler_id = 0
+
+        for arg in kwargs.keys():
+            print("warning: option %s ignored when using --libserver" % arg)
+
+        if web:
+            os.chdir(web)
+            
+        HTTPServer.__init__(self, (listen_host, listen_port), 
+                            RequestHandlerClass)
+
+
+    def process_request(self, request, client_address):
+        """Override process_request to implement a counter"""
+        self.handler_id += 1
+        ForkingMixIn.process_request(self, request, client_address)
+
 
 if __name__ == '__main__':
     websockify_init()