|
|
- import re
- import errno
- import socket
- import struct
- import collections
-
- from base64 import b64decode, b64encode
- from hashlib import sha1
-
-
- WS_KEY = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
-
-
- def pack_message(message):
- """Pack the message inside ``00`` and ``FF``
- As per the dataframing section (5.3) for the websocket spec
- """
- if isinstance(message, unicode):
- message = message.encode('utf-8')
- elif not isinstance(message, str):
- message = str(message)
- packed = "\x00%s\xFF" % message
- return packed
-
-
- def encode_hybi(buf, opcode, base64=False):
- """ Encode a HyBi style WebSocket frame.
- Optional opcode:
- 0x0 - continuation
- 0x1 - text frame (base64 encode buf)
- 0x2 - binary frame (use raw buf)
- 0x8 - connection close
- 0x9 - ping
- 0xA - pong
- """
- if base64:
- buf = b64encode(buf)
- b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
- payload_len = len(buf)
- if payload_len <= 125:
- header = struct.pack('>BB', b1, payload_len)
- elif payload_len > 125 and payload_len < 65536:
- header = struct.pack('>BBH', b1, 126, payload_len)
- elif payload_len >= 65536:
- header = struct.pack('>BBQ', b1, 127, payload_len)
- return header + buf, len(header), 0
-
-
- def decode_hybi(buf, base64=False):
- """Decode HyBi style WebSocket packets."""
- f = {'fin' : 0,
- 'opcode' : 0,
- 'mask' : 0,
- 'hlen' : 2,
- 'length' : 0,
- 'payload' : None,
- 'left' : 0,
- 'close_code' : None,
- 'close_reason' : None}
-
- blen = len(buf)
- f['left'] = blen
-
- if blen < f['hlen']:
- return f # Incomplete frame header
-
- b1, b2 = struct.unpack_from(">BB", buf)
- f['opcode'] = b1 & 0x0f
- f['fin'] = (b1 & 0x80) >> 7
- has_mask = (b2 & 0x80) >> 7
-
- f['length'] = b2 & 0x7f
-
- if f['length'] == 126:
- f['hlen'] = 4
- if blen < f['hlen']:
- return f # Incomplete frame header
- (f['length'],) = struct.unpack_from('>xxH', buf)
- elif f['length'] == 127:
- f['hlen'] = 10
- if blen < f['hlen']:
- return f # Incomplete frame header
- (f['length'],) = struct.unpack_from('>xxQ', buf)
-
- full_len = f['hlen'] + has_mask * 4 + f['length']
-
- if blen < full_len: # Incomplete frame
- return f # Incomplete frame header
-
- # Number of bytes that are part of the next frame(s)
- f['left'] = blen - full_len
-
- # Process 1 frame
- if has_mask:
- # unmask payload
- f['mask'] = buf[f['hlen']:f['hlen']+4]
- b = c = ''
- if f['length'] >= 4:
- data = struct.unpack('<I', buf[f['hlen']:f['hlen']+4])[0]
- of1 = f['hlen']+4
- b = ''
- for i in xrange(0, int(f['length']/4)):
- mask = struct.unpack('<I', buf[of1+4*i:of1+4*(i+1)])[0]
- b += struct.pack('I', data ^ mask)
-
- if f['length'] % 4:
- l = f['length'] % 4
- of1 = f['hlen']
- of2 = full_len - l
- c = ''
- for i in range(0, l):
- mask = struct.unpack('B', buf[of1 + i])[0]
- data = struct.unpack('B', buf[of2 + i])[0]
- c += chr(data ^ mask)
-
- f['payload'] = b + c
- else:
- f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
-
- if base64 and f['opcode'] in [1, 2]:
- f['payload'] = b64decode(f['payload'])
-
- if f['opcode'] == 0x08:
- if f['length'] >= 2:
- f['close_code'] = struct.unpack_from(">H", f['payload'])
- if f['length'] > 3:
- f['close_reason'] = f['payload'][2:]
-
- return f
-
-
- class WebSocketWSGI(object):
-
- def __init__(self, handler):
- self.handler = handler
-
- def verify_client(self, ws):
- pass
-
- def __call__(self, environ, start_response):
- if not (environ.get('HTTP_CONNECTION').find('Upgrade') != -1 and
- environ['HTTP_UPGRADE'].lower() == 'websocket'):
- # need to check a few more things here for true compliance
- start_response('400 Bad Request', [('Connection','close')])
- return []
-
- sock = environ['gunicorn.socket']
-
- ws = WebSocket(sock, environ)
-
- handshake_reply = ("HTTP/1.1 101 Switching Protocols\r\n"
- "Upgrade: websocket\r\n"
- "Connection: Upgrade\r\n")
-
- path = environ['PATH_INFO']
- key = environ.get('HTTP_SEC_WEBSOCKET_KEY')
- if key:
- ws_key = b64decode(key)
- if len(ws_key) != 16:
- start_response('400 Bad Request', [('Connection','close')])
- return []
-
- protocols = []
- subprotocols = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL')
- ws_protocols = []
- if subprotocols:
- for s in subprotocols.split(','):
- s = s.strip()
- if s in protocols:
- ws_protocols.append(s)
- if ws_protocols:
- handshake_reply += 'Sec-WebSocket-Protocol: %s\r\n' % ', '.join(ws_protocols)
-
- exts = []
- extensions = environ.get('HTTP_SEC_WEBSOCKET_EXTENSIONS')
- ws_extensions = []
- if extensions:
- for ext in extensions.split(','):
- ext = ext.strip()
- if ext in exts:
- ws_extensions.append(ext)
- if ws_extensions:
- handshake_reply += 'Sec-WebSocket-Extensions: %s\r\n' % ', '.join(ws_extensions)
-
- handshake_reply += (
- "Sec-WebSocket-Origin: %s\r\n"
- "Sec-WebSocket-Location: ws://%s%s\r\n"
- "Sec-WebSocket-Version: %s\r\n"
- "Sec-WebSocket-Accept: %s\r\n\r\n"
- % (
- environ.get('HTTP_ORIGIN'),
- environ.get('HTTP_HOST'),
- path,
- ws.version,
- b64encode(sha1(key + WS_KEY).digest())
- ))
-
- else:
-
- handshake_reply += (
- "WebSocket-Origin: %s\r\n"
- "WebSocket-Location: ws://%s%s\r\n\r\n" % (
- environ.get('HTTP_ORIGIN'),
- environ.get('HTTP_HOST'),
- path))
-
- sock.sendall(handshake_reply)
-
- try:
- self.handler(ws)
- except socket.error as e:
- if e[0] != errno.EPIPE:
- raise
-
- # use this undocumented feature of grainbows to ensure that it
- # doesn't barf on the fact that we didn't call start_response
- return ALREADY_HANDLED
-
-
- class WebSocket(object):
-
- def __init__(self, sock, environ, version=76):
- self._socket = sock
- try:
- version = int(environ.get('HTTP_SEC_WEBSOCKET_VERSION'))
- except (ValueError, TypeError):
- version = 76
- self.version = version
- self.closed = False
- self.accepted = False
- self._buf = b''
- self._msgs = collections.deque()
-
- def _parse_messages(self):
- """ Parses for messages in the buffer *buf*. It is assumed that
- the buffer contains the start character for a message, but that it
- may contain only part of the rest of the message.
- Returns an array of messages, and the buffer remainder that
- didn't contain any full messages."""
- msgs = []
- end_idx = 0
- buf = self._buf
- while buf:
- if self.version in (7, 8, 13):
- frame = decode_hybi(buf, base64=False)
-
- if frame['payload'] == None:
- break
- else:
- if frame['opcode'] == 0x8: # connection close
- self.closed = True
- break
- else:
- msgs.append(frame['payload']);
- if frame['left']:
- buf = buf[-frame['left']:]
- else:
- buf = b''
-
- else:
- frame_type = ord(buf[0])
- if frame_type == 0:
- # Normal message.
- end_idx = buf.find("\xFF")
- if end_idx == -1: #pragma NO COVER
- break
- msgs.append(buf[1:end_idx].decode('utf-8', 'replace'))
- buf = buf[end_idx+1:]
- elif frame_type == 255:
- # Closing handshake.
- assert ord(buf[1]) == 0, "Unexpected closing handshake: %r" % buf
- self.closed = True
- break
- else:
- raise ValueError("Don't understand how to parse this type of message: %r" % buf)
- self._buf = buf
- return msgs
-
- def send(self, message):
- """Send a message to the browser.
- *message* should be convertable to a string; unicode objects should be
- encodable as utf-8. Raises socket.error with errno of 32
- (broken pipe) if the socket has already been closed by the client.
- """
- if self.version in (7, 8, 13):
- packed, lenhead, lentail = encode_hybi(
- message, opcode=0x01, base64=False)
- else:
- packed = pack_message(message)
- self._socket.sendall(packed)
-
- def wait(self):
- """Waits for and deserializes messages.
- Returns a single message; the oldest not yet processed. If the client
- has already closed the connection, returns None. This is different
- from normal socket behavior because the empty string is a valid
- websocket message."""
- while not self._msgs:
- # Websocket might be closed already.
- if self.closed:
- return None
- # no parsed messages, must mean buf needs more data
- delta = self._socket.recv(8096)
- if delta == '':
- return None
- self._buf += delta
- msgs = self._parse_messages()
- self._msgs.extend(msgs)
- return self._msgs.popleft()
-
- def _send_closing_frame(self, ignore_send_errors=False):
- """Sends the closing frame to the client, if required."""
- if self.version in (7, 8, 13) and not self.closed:
- msg = ''
- #if code != None:
- # msg = struct.pack(">H%ds" % (len(reason)), code)
-
- buf, h, t = encode_hybi(msg, opcode=0x08, base64=False)
- self._socket.sendall(buf)
- self.closed = True
-
- elif self.version == 76 and not self.closed:
- try:
- self._socket.sendall("\xff\x00")
- except socket.error:
- # Sometimes, like when the remote side cuts off the connection,
- # we don't care about this.
- if not ignore_send_errors: #pragma NO COVER
- raise
- self.closed = True
-
- def close(self):
- """Forcibly close the websocket; generally it is preferable to
- return from the handler method."""
- self._send_closing_frame()
- self._socket.shutdown(True)
- self._socket.close()
|