From 5d2c38640015ad09c4073937ff15dd8c15aef3bc Mon Sep 17 00:00:00 2001
From: Joel Martin <github@martintribe.org>
Date: Fri, 16 Apr 2010 16:34:19 -0500
Subject: [PATCH] Reassemble partial client packets in wsproxy.py

---
 vnc.js     | 34 ++++++++++++++++++++--------------
 wsproxy.py | 35 +++++++++++++++++++++++++++++++----
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/vnc.js b/vnc.js
index e5f6657..c9704db 100644
--- a/vnc.js
+++ b/vnc.js
@@ -802,33 +802,39 @@ init_ws: function () {
     console.log("connecting to " + uri);
     RFB.ws = new WebSocket(uri);
     RFB.ws.onmessage = function(e) {
-        //console.log(">> onmessage");
+        //console.log(">> WebSockets.onmessage");
         RFB.d = RFB.d.concat(Base64.decode(e.data));
-        if (RFB.state != 'normal') {
-            RFB.init_msg();
-        } else {
-            RFB.normal_msg();
-        }
-        if (RFB.state == 'reset') {
+        if (RFB.state == 'closed') {
+            console.log("onmessage while close");
+        } else if (RFB.state == 'reset') {
             /* close and reset connection */
             RFB.disconnect();
             RFB.init_ws();
         } else if (RFB.state == 'failed') {
             console.log("Giving up!");
             RFB.disconnect();
+        } else if (RFB.state != 'normal') {
+            RFB.init_msg();
+        } else {
+            RFB.normal_msg();
         }
-        //console.log("<< onmessage");
+        //console.log("<< WebSockets.onmessage");
     };
     RFB.ws.onopen = function(e) {
-        console.log(">> onopen");
+        console.log(">> WebSockets.onopen");
         RFB.state = "ProtocolVersion";
-        console.log("<< onopen");
+        console.log("<< WebSockets.onopen");
     };
     RFB.ws.onclose = function(e) {
-        console.log(">> onclose");
+        console.log(">> WebSockets.onclose");
         RFB.state = "closed";
-        console.log("<< onclose");
-    }
+        console.log("<< WebSockets.onclose");
+    };
+    RFB.ws.onerror = function(e) {
+        console.log(">> WebSockets.onerror");
+        console.log("   " + e);
+        console.log("<< WebSockets.onerror");
+    };
 
     console.log("<< init_ws");
 },
@@ -838,7 +844,7 @@ connect: function () {
     RFB.host = $('host').value;
     RFB.port = $('port').value;
     RFB.password = $('password').value;
-    if ((!host) || (!port)) {
+    if ((!RFB.host) || (!RFB.port)) {
         console.log("must set host and port");
         return;
     }
diff --git a/wsproxy.py b/wsproxy.py
index 8653804..13d3f81 100755
--- a/wsproxy.py
+++ b/wsproxy.py
@@ -15,8 +15,16 @@ WebSocket-Protocol: sample\r
 \r
 """
 
+policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>"""
+
 def handshake(client):
-    handshake = client.recv(255)
+    handshake = client.recv(1024)
+    print "Handshake [%s]" % handshake
+    if handshake.startswith("<policy-file-request/>"):
+        print "Sending:", policy_response
+        client.send(policy_response)
+        handshake = client.recv(1024)
+        print "Handshake [%s]" % handshake
     req_lines = handshake.split("\r\n")
     _, path, _ = req_lines[0].split(" ")
     _, origin = req_lines[4].split(" ")
@@ -27,8 +35,17 @@ def traffic(token="."):
     sys.stdout.write(token)
     sys.stdout.flush()
 
+def decode(buf):
+    """ Parse out WebSocket packets. """
+    if buf.count('\xff') > 1:
+        return [d[1:] for d in buf.split('\xff')]
+    else:
+        return [b64decode(buf[1:-1])]
+
 def proxy(client, target):
+    """ Proxy WebSocket to normal socket. """
     cqueue = []
+    cpartial = ""
     tqueue = []
     socks = [client, target]
 
@@ -39,9 +56,19 @@ def proxy(client, target):
         if client in ins:
             buf = client.recv(buffer_size)
             if len(buf) == 0: raise Exception("Client closed")
-            tqueue.append(b64decode(buf[1:-1]))
-            #print "Client recv: %s (%d)" % (repr(buf[1:-1]), len(buf))
-            traffic("}")
+
+            if buf[-1] == "\xff":
+                if cpartial:
+                    tqueue.extend(decode(cpartial + buf))
+                    cpartial = ""
+                else:
+                    tqueue.extend(decode(buf))
+                traffic("}")
+            else:
+                traffic(".}")
+                cpartial = cpartial + buf
+
+            #print "Client recv: %s (%d)" % (repr(buf, len(buf))
 
         if target in ins:
             buf = target.recv(buffer_size)