sync with websockify

Pull in websockify 1669139

3 files changed, 52 insertions, 30 deletions

diff --git a/include/websock.js b/include/websock.js
index 9d6a306..7d12644 100644
--- a/include/websock.js
+++ b/include/websock.js
@@ -172,7 +172,10 @@ function decode_message(data) {
     //Util.Debug(">> decode_message: " + data);
     if (mode === 'binary') {
         // push arraybuffer values onto the end
-        rQ.push.apply(rQ, (new Uint8Array(data)));
+        var u8 = new Uint8Array(data);
+        for (var i = 0; i < u8.length; i++) {
+            rQ.push(u8[i]);
+        }
     } else {
         // base64 decode and concat to the end
         rQ = rQ.concat(Base64.decode(data, 0));
diff --git a/utils/websocket.py b/utils/websocket.py
index 4c0cacc..1a5b9ff 100644
--- a/utils/websocket.py
+++ b/utils/websocket.py
@@ -240,32 +240,33 @@ Sec-WebSocket-Accept: %s\r
         os.dup2(os.open(os.devnull, os.O_RDWR), sys.stderr.fileno())
 
     @staticmethod
-    def unmask(buf, f):
-        pstart = f['hlen'] + 4
-        pend = pstart + f['length']
+    def unmask(buf, hlen, plen):
+        pstart = hlen + 4
+        pend = pstart + plen
         if numpy:
             b = c = s2b('')
-            if f['length'] >= 4:
+            if plen >= 4:
                 mask = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
-                        offset=f['hlen'], count=1)
+                        offset=hlen, count=1)
                 data = numpy.frombuffer(buf, dtype=numpy.dtype('<u4'),
-                        offset=pstart, count=int(f['length'] / 4))
+                        offset=pstart, count=int(plen / 4))
                 #b = numpy.bitwise_xor(data, mask).data
                 b = numpy.bitwise_xor(data, mask).tostring()
 
-            if f['length'] % 4:
+            if plen % 4:
                 #print("Partial unmask")
                 mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
-                        offset=f['hlen'], count=(f['length'] % 4))
+                        offset=hlen, count=(plen % 4))
                 data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
-                        offset=pend - (f['length'] % 4),
-                        count=(f['length'] % 4))
+                        offset=pend - (plen % 4),
+                        count=(plen % 4))
                 c = numpy.bitwise_xor(data, mask).tostring()
             return b + c
         else:
             # Slower fallback
+            mask = buf[hlen:hlen+4]
             data = array.array('B')
-            mask = s2a(f['mask'])
+            mask = s2a(mask)
             data.fromstring(buf[pstart:pend])
             for i in range(len(data)):
                 data[i] ^= mask[i % 4]
@@ -304,7 +305,7 @@ Sec-WebSocket-Accept: %s\r
         Returns:
             {'fin'          : 0_or_1,
              'opcode'       : number,
-             'mask'         : 32_bit_number,
+             'masked'       : boolean,
              'hlen'         : header_bytes_number,
              'length'       : payload_bytes_number,
              'payload'      : decoded_buffer,
@@ -315,7 +316,7 @@ Sec-WebSocket-Accept: %s\r
 
         f = {'fin'          : 0,
              'opcode'       : 0,
-             'mask'         : 0,
+             'masked'       : False,
              'hlen'         : 2,
              'length'       : 0,
              'payload'      : None,
@@ -332,7 +333,7 @@ Sec-WebSocket-Accept: %s\r
         b1, b2 = unpack_from(">BB", buf)
         f['opcode'] = b1 & 0x0f
         f['fin'] = (b1 & 0x80) >> 7
-        has_mask = (b2 & 0x80) >> 7
+        f['masked'] = (b2 & 0x80) >> 7
 
         f['length'] = b2 & 0x7f
 
@@ -347,7 +348,7 @@ Sec-WebSocket-Accept: %s\r
                 return f # Incomplete frame header
             (f['length'],) = unpack_from('>xxQ', buf)
 
-        full_len = f['hlen'] + has_mask * 4 + f['length']
+        full_len = f['hlen'] + f['masked'] * 4 + f['length']
 
         if blen < full_len: # Incomplete frame
             return f # Incomplete frame header
@@ -356,13 +357,13 @@ Sec-WebSocket-Accept: %s\r
         f['left'] = blen - full_len
 
         # Process 1 frame
-        if has_mask:
+        if f['masked']:
             # unmask payload
-            f['mask'] = buf[f['hlen']:f['hlen']+4]
-            f['payload'] = WebSocketServer.unmask(buf, f)
+            f['payload'] = WebSocketServer.unmask(buf, f['hlen'],
+                                                  f['length'])
         else:
             print("Unmasked frame: %s" % repr(buf))
-            f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
+            f['payload'] = buf[(f['hlen'] + f['masked'] * 4):full_len]
 
         if base64 and f['opcode'] in [1, 2]:
             try:
@@ -389,6 +390,7 @@ Sec-WebSocket-Accept: %s\r
         end = buf.find(s2b('\xff'))
         return {'payload': b64decode(buf[1:end]),
                 'hlen': 1,
+                'masked': False,
                 'length': end - 1,
                 'left': len(buf) - (end + 1)}
 
@@ -456,7 +458,7 @@ Sec-WebSocket-Accept: %s\r
                 if self.rec:
                     self.rec.write("%s,\n" %
                             repr("{%s{" % tdelta
-                                + encbuf[lenhead:-lentail]))
+                                + encbuf[lenhead:len(encbuf)-lentail]))
 
                 self.send_parts.append(encbuf)
 
@@ -536,8 +538,14 @@ Sec-WebSocket-Accept: %s\r
             if self.rec:
                 start = frame['hlen']
                 end = frame['hlen'] + frame['length']
+                if frame['masked']:
+                    recbuf = WebSocketServer.unmask(buf, frame['hlen'],
+                                                   frame['length'])
+                else:
+                    recbuf = buf[frame['hlen']:frame['hlen'] +
+                                               frame['length']]
                 self.rec.write("%s,\n" %
-                        repr("}%s}" % tdelta + buf[start:end]))
+                        repr("}%s}" % tdelta + recbuf))
 
 
             bufs.append(frame['payload'])
@@ -779,6 +787,10 @@ Sec-WebSocket-Accept: %s\r
                                         self.handler_id)
                     self.msg("opening record file: %s" % fname)
                     self.rec = open(fname, 'w+')
+                    encoding = "binary"
+                    if self.base64: encoding = "base64"
+                    self.rec.write("var VNC_frame_encoding = '%s';\n"
+                            % encoding)
                     self.rec.write("var VNC_frame_data = [\n")
 
                 self.ws_connection = True
@@ -800,7 +812,7 @@ Sec-WebSocket-Accept: %s\r
                     self.msg(traceback.format_exc())
         finally:
             if self.rec:
-                self.rec.write("'EOF']\n")
+                self.rec.write("'EOF'];\n")
                 self.rec.close()
 
             if self.client and self.client != startsock:
diff --git a/utils/websockify b/utils/websockify
index c792dbf..8bd67ec 100755
--- a/utils/websockify
+++ b/utils/websockify
@@ -11,7 +11,7 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
 
 '''
 
-import socket, optparse, time, os, sys, subprocess
+import signal, socket, optparse, time, os, sys, subprocess
 from select import select
 import websocket
 try:    from urllib.parse import parse_qs, urlparse
@@ -87,7 +87,7 @@ Traffic Legend:
         self.wrap_times.append(time.time())
         self.wrap_times.pop(0)
         self.cmd = subprocess.Popen(
-                self.wrap_cmd, env=os.environ)
+                self.wrap_cmd, env=os.environ, preexec_fn=_subprocess_setup)
         self.spawn_message = True
 
     def started(self):
@@ -144,7 +144,7 @@ Traffic Legend:
                 else:
                     self.run_wrap_cmd()
 
-    # 
+    #
     # Routines above this point are run in the master listener
     # process.
     #
@@ -165,13 +165,13 @@ Traffic Legend:
 
         # Connect to the target
         if self.wrap_cmd:
-            msg = "connecting to command: %s" % (" ".join(self.wrap_cmd), self.target_port)
+            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)
@@ -291,6 +291,13 @@ Traffic Legend:
                         self.target_host, self.target_port))
                     raise self.CClose(closed['code'], closed['reason'])
 
+
+def _subprocess_setup():
+    # Python installs a SIGPIPE handler by default. This is usually not what
+    # non-Python successfulbprocesses expect.
+    signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+
+
 def websockify_init():
     usage = "\n    %prog [options]"
     usage += " [source_addr:]source_port [target_addr:target_port]"
@@ -338,7 +345,7 @@ def websockify_init():
     (opts, args) = parser.parse_args()
 
     # Sanity checks
-    if len(args) < 2 and not opts.target_cfg:
+    if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
         parser.error("Too few arguments")
     if sys.argv.count('--'):
         opts.wrap_cmd = args[1:]
@@ -349,7 +356,7 @@ def websockify_init():
 
     if not websocket.ssl and opts.ssl_target:
         parser.error("SSL target requested and Python SSL module not loaded.");
-        
+
     if opts.ssl_only and not os.path.exists(opts.cert):
         parser.error("SSL only and %s not found" % opts.cert)