Add a --run-once option to wsproxy/websockify

[krims0n32]

In my project I run wsproxy only when needed (e.g. a VNC session is requested by the user). It would be nice if wsproxy had a --run-once option that would listen for the incoming connection, setup the session, and once the user disconnects (or any connection error occurs) would terminate itself.

Currently I did an ugly hack in the wsproxy script to accomplish this by adding a few extra exceptions and handling accordingly (restart when a flash policy request is sent, terminate on disconnect or error), but there must be a better way of achieving this :)

I need to add that I know you can run a process and have wsproxy terminate when that process exits but well then I would need to put in an extra process to tunnel between VNC (KVM) and wsproxy, eating up unnecessary CPU.

Thanks for this project, it beats the java based solution I used before in many ways.

[kanaka]

Can you try the following patch and see if it works for you. It won't exit on handshake errors, but if the handshake is a valid WebSocket connection upgrade, then it will exit after that (regardless of if the WebSocket connections terminates normally or not).

diff --git a/utils/websocket.py b/utils/websocket.py
index 8194914..c84f885 100644
--- a/utils/websocket.py
+++ b/utils/websocket.py
@@ -88,7 +88,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=''):
+            daemon=False, run_once=False, record='', web=''):

         # settings
         self.verbose        = verbose
@@ -96,6 +96,7 @@ Sec-WebSocket-Accept: %s\r
         self.listen_port    = listen_port
         self.ssl_only       = ssl_only
         self.daemon         = daemon
+        self.run_once       = run_once
         self.handler_id     = 1

         # Make paths settings absolute
@@ -714,6 +715,8 @@ Sec-WebSocket-Accept: %s\r
         self.rec        = None
         self.start_time = int(time.time()*1000)

+        ws_connection = False
+
         # handler process
         try:
             try:
@@ -727,6 +730,7 @@ Sec-WebSocket-Accept: %s\r
                     self.rec = open(fname, 'w+')
                     self.rec.write("var VNC_frame_data = [\n")

+                ws_connection = True
                 self.new_client()
             except self.EClose:
                 _, exc, _ = sys.exc_info()
@@ -746,6 +750,8 @@ Sec-WebSocket-Accept: %s\r
             if self.client and self.client != startsock:
                 self.client.close()

+        return ws_connection
+
     def new_client(self):
         """ Do something with a WebSockets client connection. """
         raise("WebSocketServer.new_client() must be overloaded")
@@ -799,7 +805,13 @@ Sec-WebSocket-Accept: %s\r
                         else:
                             raise

-                    if Process:
+                    if self.run_once:
+                        ret = self.top_new_client(startsock, address)
+                        if ret:
+                            self.vmsg('%s: exiting after connection'
+                                    % address[0])
+                            break
+                    elif Process:
                         self.vmsg('%s: new handler Process' % address[0])
                         p = Process(target=self.top_new_client,
                                 args=(startsock, address))
diff --git a/utils/websockify b/utils/websockify
index 59b6e0d..468c0f4 100755
--- a/utils/websockify
+++ b/utils/websockify
@@ -226,6 +226,8 @@ if __name__ == '__main__':
     parser.add_option("--daemon", "-D",
             dest="daemon", action="store_true",
             help="become a daemon (background process)")
+    parser.add_option("--run-once", action="store_true",
+            help="handle a single WebSocket connection and exit")
     parser.add_option("--cert", default="self.pem",
             help="SSL certificate file")
     parser.add_option("--key", default=None,

[krims0n32]

Thanks, I did some simple testing with Firefox and this seems to do what I want :) I will give it some more thorough testing using different browsers. Perhaps it is better to also terminate when any error occurs (including handshake). I want to prevent that websockify gets "stuck" waiting for nothing (e.g. the client aborted, a network error occured, etc.). I forgot to add I also hacked in a select timeout of 15 seconds, if nothing comes in during that period, websockify assumes a client problem and terminates as well. That way there will be no ghost websockify processes with no clients interacting with them. If this happens I raise a "Timed out" exception that in turn writes the error to stderr so my parent process can report this error back to the client. Would it be possible to put that in as well ? Example below:

try:
    self.poll()
    ready = select.select([lsock], [], [], 15)[0]
    if lsock in ready:
        startsock, address = lsock.accept()
    else:
        raise Exception('Timed out')

...

except Exception:
    _, exc, _ = sys.exc_info()
    self.msg("handler exception: %s" % str(exc))
    sys.stderr.write(str(exc))
    if self.verbose:
        self.msg(traceback.format_exc())
    break

Thanks for your work.

[kanaka]

The challenge is that the solution needs to work well with the flash policy response mechanism and the builtin webserver.

Try this and see if it works for you: https://gist.github.com/1229802 It's a slight change to the --run-once patch and also adds a --timeout TIMEOUT option

Let me know you results and I'll push it to out after that.

[krims0n32]

This seems to work rather well ! Only thing I need now is for my parent process (which launches wsproxy) to know what happened. Could you perhaps add a returncode so I can tell something went wrong (timeout or some other error) ?

[kanaka]

@krims0n32, I pushed out the --run-once and --timeout changes to websockify (I probably won't merge to noVNC until there are more substantial changes). I don't have the time right now to work on return codes (it's not quite as trivial as --run-once or --timeout). However, it seems like useful functionality so if you come up with a clean solution then send me a pull request. Closing this bug since it is resolved.