From 636aba62ed68ea4907646205ed2eb6e8f24fa1ba Mon Sep 17 00:00:00 2001
From: Joel Martin <github@martintribe.org>
Date: Thu, 22 Sep 2011 15:52:02 -0500
Subject: [PATCH] Add --run-once and --timeout TIME parameters.

- --run-once will exit after handling a single WebSocket connection
  (but not ater flash policy or normal web requests).

- --timeout TIME will stop listening for new connections after exit
  after TIME seconds (the master process shuts down). Existing
  WebSocket connections will continue but once all connections are
  closed all processes will terminate.
---
 websocket.py | 24 ++++++++++++++++++++++--
 websockify   |  4 ++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/websocket.py b/websocket.py
index 8194914..833d2a2 100644
--- a/websocket.py
+++ b/websocket.py
@@ -88,7 +88,8 @@ 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, record='', web='',
+            run_once=False, timeout=0):
 
         # settings
         self.verbose        = verbose
@@ -96,6 +97,11 @@ Sec-WebSocket-Accept: %s\r
         self.listen_port    = listen_port
         self.ssl_only       = ssl_only
         self.daemon         = daemon
+        self.run_once       = run_once
+        self.timeout        = timeout
+
+        self.launch_time    = time.time()
+        self.ws_connection  = False
         self.handler_id     = 1
 
         # Make paths settings absolute
@@ -727,6 +733,7 @@ Sec-WebSocket-Accept: %s\r
                     self.rec = open(fname, 'w+')
                     self.rec.write("var VNC_frame_data = [\n")
 
+                self.ws_connection = True
                 self.new_client()
             except self.EClose:
                 _, exc, _ = sys.exc_info()
@@ -777,6 +784,12 @@ Sec-WebSocket-Accept: %s\r
                     startsock = None
                     pid = err = 0
 
+                    time_elapsed = time.time() - self.launch_time
+                    if self.timeout and time_elapsed > self.timeout:
+                        self.msg('listener exit due to --timeout %s'
+                                % self.timeout)
+                        break
+
                     try:
                         self.poll()
 
@@ -799,7 +812,14 @@ Sec-WebSocket-Accept: %s\r
                         else:
                             raise
 
-                    if Process:
+                    if self.run_once:
+                        # Run in same process if run_once
+                        self.top_new_client(startsock, address)
+                        if self.ws_connection :
+                            self.msg('%s: exiting due to --run-once'
+                                    % 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/websockify b/websockify
index 59b6e0d..4521a96 100755
--- a/websockify
+++ b/websockify
@@ -226,6 +226,10 @@ 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("--timeout", type=int, default=0,
+            help="after TIMEOUT seconds exit when not connected")
     parser.add_option("--cert", default="self.pem",
             help="SSL certificate file")
     parser.add_option("--key", default=None,