You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

337 lines
11 KiB

4 years ago
  1. import re
  2. import errno
  3. import socket
  4. import struct
  5. import collections
  6. from base64 import b64decode, b64encode
  7. from hashlib import sha1
  8. WS_KEY = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
  9. def pack_message(message):
  10. """Pack the message inside ``00`` and ``FF``
  11. As per the dataframing section (5.3) for the websocket spec
  12. """
  13. if isinstance(message, unicode):
  14. message = message.encode('utf-8')
  15. elif not isinstance(message, str):
  16. message = str(message)
  17. packed = "\x00%s\xFF" % message
  18. return packed
  19. def encode_hybi(buf, opcode, base64=False):
  20. """ Encode a HyBi style WebSocket frame.
  21. Optional opcode:
  22. 0x0 - continuation
  23. 0x1 - text frame (base64 encode buf)
  24. 0x2 - binary frame (use raw buf)
  25. 0x8 - connection close
  26. 0x9 - ping
  27. 0xA - pong
  28. """
  29. if base64:
  30. buf = b64encode(buf)
  31. b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
  32. payload_len = len(buf)
  33. if payload_len <= 125:
  34. header = struct.pack('>BB', b1, payload_len)
  35. elif payload_len > 125 and payload_len < 65536:
  36. header = struct.pack('>BBH', b1, 126, payload_len)
  37. elif payload_len >= 65536:
  38. header = struct.pack('>BBQ', b1, 127, payload_len)
  39. return header + buf, len(header), 0
  40. def decode_hybi(buf, base64=False):
  41. """Decode HyBi style WebSocket packets."""
  42. f = {'fin' : 0,
  43. 'opcode' : 0,
  44. 'mask' : 0,
  45. 'hlen' : 2,
  46. 'length' : 0,
  47. 'payload' : None,
  48. 'left' : 0,
  49. 'close_code' : None,
  50. 'close_reason' : None}
  51. blen = len(buf)
  52. f['left'] = blen
  53. if blen < f['hlen']:
  54. return f # Incomplete frame header
  55. b1, b2 = struct.unpack_from(">BB", buf)
  56. f['opcode'] = b1 & 0x0f
  57. f['fin'] = (b1 & 0x80) >> 7
  58. has_mask = (b2 & 0x80) >> 7
  59. f['length'] = b2 & 0x7f
  60. if f['length'] == 126:
  61. f['hlen'] = 4
  62. if blen < f['hlen']:
  63. return f # Incomplete frame header
  64. (f['length'],) = struct.unpack_from('>xxH', buf)
  65. elif f['length'] == 127:
  66. f['hlen'] = 10
  67. if blen < f['hlen']:
  68. return f # Incomplete frame header
  69. (f['length'],) = struct.unpack_from('>xxQ', buf)
  70. full_len = f['hlen'] + has_mask * 4 + f['length']
  71. if blen < full_len: # Incomplete frame
  72. return f # Incomplete frame header
  73. # Number of bytes that are part of the next frame(s)
  74. f['left'] = blen - full_len
  75. # Process 1 frame
  76. if has_mask:
  77. # unmask payload
  78. f['mask'] = buf[f['hlen']:f['hlen']+4]
  79. b = c = ''
  80. if f['length'] >= 4:
  81. data = struct.unpack('<I', buf[f['hlen']:f['hlen']+4])[0]
  82. of1 = f['hlen']+4
  83. b = ''
  84. for i in xrange(0, int(f['length']/4)):
  85. mask = struct.unpack('<I', buf[of1+4*i:of1+4*(i+1)])[0]
  86. b += struct.pack('I', data ^ mask)
  87. if f['length'] % 4:
  88. l = f['length'] % 4
  89. of1 = f['hlen']
  90. of2 = full_len - l
  91. c = ''
  92. for i in range(0, l):
  93. mask = struct.unpack('B', buf[of1 + i])[0]
  94. data = struct.unpack('B', buf[of2 + i])[0]
  95. c += chr(data ^ mask)
  96. f['payload'] = b + c
  97. else:
  98. f['payload'] = buf[(f['hlen'] + has_mask * 4):full_len]
  99. if base64 and f['opcode'] in [1, 2]:
  100. f['payload'] = b64decode(f['payload'])
  101. if f['opcode'] == 0x08:
  102. if f['length'] >= 2:
  103. f['close_code'] = struct.unpack_from(">H", f['payload'])
  104. if f['length'] > 3:
  105. f['close_reason'] = f['payload'][2:]
  106. return f
  107. class WebSocketWSGI(object):
  108. def __init__(self, handler):
  109. self.handler = handler
  110. def verify_client(self, ws):
  111. pass
  112. def __call__(self, environ, start_response):
  113. if not (environ.get('HTTP_CONNECTION').find('Upgrade') != -1 and
  114. environ['HTTP_UPGRADE'].lower() == 'websocket'):
  115. # need to check a few more things here for true compliance
  116. start_response('400 Bad Request', [('Connection','close')])
  117. return []
  118. sock = environ['gunicorn.socket']
  119. ws = WebSocket(sock, environ)
  120. handshake_reply = ("HTTP/1.1 101 Switching Protocols\r\n"
  121. "Upgrade: websocket\r\n"
  122. "Connection: Upgrade\r\n")
  123. path = environ['PATH_INFO']
  124. key = environ.get('HTTP_SEC_WEBSOCKET_KEY')
  125. if key:
  126. ws_key = b64decode(key)
  127. if len(ws_key) != 16:
  128. start_response('400 Bad Request', [('Connection','close')])
  129. return []
  130. protocols = []
  131. subprotocols = environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL')
  132. ws_protocols = []
  133. if subprotocols:
  134. for s in subprotocols.split(','):
  135. s = s.strip()
  136. if s in protocols:
  137. ws_protocols.append(s)
  138. if ws_protocols:
  139. handshake_reply += 'Sec-WebSocket-Protocol: %s\r\n' % ', '.join(ws_protocols)
  140. exts = []
  141. extensions = environ.get('HTTP_SEC_WEBSOCKET_EXTENSIONS')
  142. ws_extensions = []
  143. if extensions:
  144. for ext in extensions.split(','):
  145. ext = ext.strip()
  146. if ext in exts:
  147. ws_extensions.append(ext)
  148. if ws_extensions:
  149. handshake_reply += 'Sec-WebSocket-Extensions: %s\r\n' % ', '.join(ws_extensions)
  150. handshake_reply += (
  151. "Sec-WebSocket-Origin: %s\r\n"
  152. "Sec-WebSocket-Location: ws://%s%s\r\n"
  153. "Sec-WebSocket-Version: %s\r\n"
  154. "Sec-WebSocket-Accept: %s\r\n\r\n"
  155. % (
  156. environ.get('HTTP_ORIGIN'),
  157. environ.get('HTTP_HOST'),
  158. path,
  159. ws.version,
  160. b64encode(sha1(key + WS_KEY).digest())
  161. ))
  162. else:
  163. handshake_reply += (
  164. "WebSocket-Origin: %s\r\n"
  165. "WebSocket-Location: ws://%s%s\r\n\r\n" % (
  166. environ.get('HTTP_ORIGIN'),
  167. environ.get('HTTP_HOST'),
  168. path))
  169. sock.sendall(handshake_reply)
  170. try:
  171. self.handler(ws)
  172. except socket.error as e:
  173. if e[0] != errno.EPIPE:
  174. raise
  175. # use this undocumented feature of grainbows to ensure that it
  176. # doesn't barf on the fact that we didn't call start_response
  177. return ALREADY_HANDLED
  178. class WebSocket(object):
  179. def __init__(self, sock, environ, version=76):
  180. self._socket = sock
  181. try:
  182. version = int(environ.get('HTTP_SEC_WEBSOCKET_VERSION'))
  183. except (ValueError, TypeError):
  184. version = 76
  185. self.version = version
  186. self.closed = False
  187. self.accepted = False
  188. self._buf = b''
  189. self._msgs = collections.deque()
  190. def _parse_messages(self):
  191. """ Parses for messages in the buffer *buf*. It is assumed that
  192. the buffer contains the start character for a message, but that it
  193. may contain only part of the rest of the message.
  194. Returns an array of messages, and the buffer remainder that
  195. didn't contain any full messages."""
  196. msgs = []
  197. end_idx = 0
  198. buf = self._buf
  199. while buf:
  200. if self.version in (7, 8, 13):
  201. frame = decode_hybi(buf, base64=False)
  202. if frame['payload'] == None:
  203. break
  204. else:
  205. if frame['opcode'] == 0x8: # connection close
  206. self.closed = True
  207. break
  208. else:
  209. msgs.append(frame['payload']);
  210. if frame['left']:
  211. buf = buf[-frame['left']:]
  212. else:
  213. buf = b''
  214. else:
  215. frame_type = ord(buf[0])
  216. if frame_type == 0:
  217. # Normal message.
  218. end_idx = buf.find("\xFF")
  219. if end_idx == -1: #pragma NO COVER
  220. break
  221. msgs.append(buf[1:end_idx].decode('utf-8', 'replace'))
  222. buf = buf[end_idx+1:]
  223. elif frame_type == 255:
  224. # Closing handshake.
  225. assert ord(buf[1]) == 0, "Unexpected closing handshake: %r" % buf
  226. self.closed = True
  227. break
  228. else:
  229. raise ValueError("Don't understand how to parse this type of message: %r" % buf)
  230. self._buf = buf
  231. return msgs
  232. def send(self, message):
  233. """Send a message to the browser.
  234. *message* should be convertable to a string; unicode objects should be
  235. encodable as utf-8. Raises socket.error with errno of 32
  236. (broken pipe) if the socket has already been closed by the client.
  237. """
  238. if self.version in (7, 8, 13):
  239. packed, lenhead, lentail = encode_hybi(
  240. message, opcode=0x01, base64=False)
  241. else:
  242. packed = pack_message(message)
  243. self._socket.sendall(packed)
  244. def wait(self):
  245. """Waits for and deserializes messages.
  246. Returns a single message; the oldest not yet processed. If the client
  247. has already closed the connection, returns None. This is different
  248. from normal socket behavior because the empty string is a valid
  249. websocket message."""
  250. while not self._msgs:
  251. # Websocket might be closed already.
  252. if self.closed:
  253. return None
  254. # no parsed messages, must mean buf needs more data
  255. delta = self._socket.recv(8096)
  256. if delta == '':
  257. return None
  258. self._buf += delta
  259. msgs = self._parse_messages()
  260. self._msgs.extend(msgs)
  261. return self._msgs.popleft()
  262. def _send_closing_frame(self, ignore_send_errors=False):
  263. """Sends the closing frame to the client, if required."""
  264. if self.version in (7, 8, 13) and not self.closed:
  265. msg = ''
  266. #if code != None:
  267. # msg = struct.pack(">H%ds" % (len(reason)), code)
  268. buf, h, t = encode_hybi(msg, opcode=0x08, base64=False)
  269. self._socket.sendall(buf)
  270. self.closed = True
  271. elif self.version == 76 and not self.closed:
  272. try:
  273. self._socket.sendall("\xff\x00")
  274. except socket.error:
  275. # Sometimes, like when the remote side cuts off the connection,
  276. # we don't care about this.
  277. if not ignore_send_errors: #pragma NO COVER
  278. raise
  279. self.closed = True
  280. def close(self):
  281. """Forcibly close the websocket; generally it is preferable to
  282. return from the handler method."""
  283. self._send_closing_frame()
  284. self._socket.shutdown(True)
  285. self._socket.close()