websocket.py 9.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267
  1. import collections
  2. import errno
  3. import string
  4. import struct
  5. from socket import error as SocketError
  6. try:
  7. from hashlib import md5
  8. except ImportError: #pragma NO COVER
  9. from md5 import md5
  10. import eventlet
  11. from eventlet import semaphore
  12. from eventlet import wsgi
  13. from eventlet.green import socket
  14. from eventlet.support import get_errno
  15. ACCEPTABLE_CLIENT_ERRORS = set((errno.ECONNRESET, errno.EPIPE))
  16. __all__ = ["WebSocketWSGI", "WebSocket"]
  17. class WebSocketWSGI(object):
  18. """Wraps a websocket handler function in a WSGI application.
  19. Use it like this::
  20. @websocket.WebSocketWSGI
  21. def my_handler(ws):
  22. from_browser = ws.wait()
  23. ws.send("from server")
  24. The single argument to the function will be an instance of
  25. :class:`WebSocket`. To close the socket, simply return from the
  26. function. Note that the server will log the websocket request at
  27. the time of closure.
  28. """
  29. def __init__(self, handler):
  30. self.handler = handler
  31. self.protocol_version = None
  32. def __call__(self, environ, start_response):
  33. if not (environ.get('HTTP_CONNECTION') == 'Upgrade' and
  34. environ.get('HTTP_UPGRADE') == 'WebSocket'):
  35. # need to check a few more things here for true compliance
  36. start_response('400 Bad Request', [('Connection','close')])
  37. return []
  38. # See if they sent the new-format headers
  39. if 'HTTP_SEC_WEBSOCKET_KEY1' in environ:
  40. self.protocol_version = 76
  41. if 'HTTP_SEC_WEBSOCKET_KEY2' not in environ:
  42. # That's bad.
  43. start_response('400 Bad Request', [('Connection','close')])
  44. return []
  45. else:
  46. self.protocol_version = 75
  47. # Get the underlying socket and wrap a WebSocket class around it
  48. sock = environ['eventlet.input'].get_socket()
  49. ws = WebSocket(sock, environ, self.protocol_version)
  50. # If it's new-version, we need to work out our challenge response
  51. if self.protocol_version == 76:
  52. key1 = self._extract_number(environ['HTTP_SEC_WEBSOCKET_KEY1'])
  53. key2 = self._extract_number(environ['HTTP_SEC_WEBSOCKET_KEY2'])
  54. # There's no content-length header in the request, but it has 8
  55. # bytes of data.
  56. environ['wsgi.input'].content_length = 8
  57. key3 = environ['wsgi.input'].read(8)
  58. key = struct.pack(">II", key1, key2) + key3
  59. response = md5(key).digest()
  60. # Start building the response
  61. scheme = 'ws'
  62. if environ.get('wsgi.url_scheme') == 'https':
  63. scheme = 'wss'
  64. location = '%s://%s%s%s' % (
  65. scheme,
  66. environ.get('HTTP_HOST'),
  67. environ.get('SCRIPT_NAME'),
  68. environ.get('PATH_INFO')
  69. )
  70. qs = environ.get('QUERY_STRING')
  71. if qs is not None:
  72. location += '?' + qs
  73. if self.protocol_version == 75:
  74. handshake_reply = ("HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
  75. "Upgrade: WebSocket\r\n"
  76. "Connection: Upgrade\r\n"
  77. "WebSocket-Origin: %s\r\n"
  78. "WebSocket-Location: %s\r\n\r\n" % (
  79. environ.get('HTTP_ORIGIN'),
  80. location))
  81. elif self.protocol_version == 76:
  82. handshake_reply = ("HTTP/1.1 101 WebSocket Protocol Handshake\r\n"
  83. "Upgrade: WebSocket\r\n"
  84. "Connection: Upgrade\r\n"
  85. "Sec-WebSocket-Origin: %s\r\n"
  86. "Sec-WebSocket-Protocol: %s\r\n"
  87. "Sec-WebSocket-Location: %s\r\n"
  88. "\r\n%s"% (
  89. environ.get('HTTP_ORIGIN'),
  90. environ.get('HTTP_SEC_WEBSOCKET_PROTOCOL', 'default'),
  91. location,
  92. response))
  93. else: #pragma NO COVER
  94. raise ValueError("Unknown WebSocket protocol version.")
  95. sock.sendall(handshake_reply)
  96. try:
  97. self.handler(ws)
  98. except socket.error, e:
  99. if get_errno(e) not in ACCEPTABLE_CLIENT_ERRORS:
  100. raise
  101. # Make sure we send the closing frame
  102. ws._send_closing_frame(True)
  103. # use this undocumented feature of eventlet.wsgi to ensure that it
  104. # doesn't barf on the fact that we didn't call start_response
  105. return wsgi.ALREADY_HANDLED
  106. def _extract_number(self, value):
  107. """
  108. Utility function which, given a string like 'g98sd 5[]221@1', will
  109. return 9852211. Used to parse the Sec-WebSocket-Key headers.
  110. """
  111. out = ""
  112. spaces = 0
  113. for char in value:
  114. if char in string.digits:
  115. out += char
  116. elif char == " ":
  117. spaces += 1
  118. return int(out) / spaces
  119. class WebSocket(object):
  120. """A websocket object that handles the details of
  121. serialization/deserialization to the socket.
  122. The primary way to interact with a :class:`WebSocket` object is to
  123. call :meth:`send` and :meth:`wait` in order to pass messages back
  124. and forth with the browser. Also available are the following
  125. properties:
  126. path
  127. The path value of the request. This is the same as the WSGI PATH_INFO variable, but more convenient.
  128. protocol
  129. The value of the Websocket-Protocol header.
  130. origin
  131. The value of the 'Origin' header.
  132. environ
  133. The full WSGI environment for this request.
  134. """
  135. def __init__(self, sock, environ, version=76):
  136. """
  137. :param socket: The eventlet socket
  138. :type socket: :class:`eventlet.greenio.GreenSocket`
  139. :param environ: The wsgi environment
  140. :param version: The WebSocket spec version to follow (default is 76)
  141. """
  142. self.socket = sock
  143. self.origin = environ.get('HTTP_ORIGIN')
  144. self.protocol = environ.get('HTTP_WEBSOCKET_PROTOCOL')
  145. self.path = environ.get('PATH_INFO')
  146. self.environ = environ
  147. self.version = version
  148. self.websocket_closed = False
  149. self._buf = ""
  150. self._msgs = collections.deque()
  151. self._sendlock = semaphore.Semaphore()
  152. @staticmethod
  153. def _pack_message(message):
  154. """Pack the message inside ``00`` and ``FF``
  155. As per the dataframing section (5.3) for the websocket spec
  156. """
  157. if isinstance(message, unicode):
  158. message = message.encode('utf-8')
  159. elif not isinstance(message, str):
  160. message = str(message)
  161. packed = "\x00%s\xFF" % message
  162. return packed
  163. def _parse_messages(self):
  164. """ Parses for messages in the buffer *buf*. It is assumed that
  165. the buffer contains the start character for a message, but that it
  166. may contain only part of the rest of the message.
  167. Returns an array of messages, and the buffer remainder that
  168. didn't contain any full messages."""
  169. msgs = []
  170. end_idx = 0
  171. buf = self._buf
  172. while buf:
  173. frame_type = ord(buf[0])
  174. if frame_type == 0:
  175. # Normal message.
  176. end_idx = buf.find("\xFF")
  177. if end_idx == -1: #pragma NO COVER
  178. break
  179. msgs.append(buf[1:end_idx].decode('utf-8', 'replace'))
  180. buf = buf[end_idx+1:]
  181. elif frame_type == 255:
  182. # Closing handshake.
  183. assert ord(buf[1]) == 0, "Unexpected closing handshake: %r" % buf
  184. self.websocket_closed = True
  185. break
  186. else:
  187. raise ValueError("Don't understand how to parse this type of message: %r" % buf)
  188. self._buf = buf
  189. return msgs
  190. def send(self, message):
  191. """Send a message to the browser.
  192. *message* should be convertable to a string; unicode objects should be
  193. encodable as utf-8. Raises socket.error with errno of 32
  194. (broken pipe) if the socket has already been closed by the client."""
  195. packed = self._pack_message(message)
  196. # if two greenthreads are trying to send at the same time
  197. # on the same socket, sendlock prevents interleaving and corruption
  198. self._sendlock.acquire()
  199. try:
  200. self.socket.sendall(packed)
  201. finally:
  202. self._sendlock.release()
  203. def wait(self):
  204. """Waits for and deserializes messages.
  205. Returns a single message; the oldest not yet processed. If the client
  206. has already closed the connection, returns None. This is different
  207. from normal socket behavior because the empty string is a valid
  208. websocket message."""
  209. while not self._msgs:
  210. # Websocket might be closed already.
  211. if self.websocket_closed:
  212. return None
  213. # no parsed messages, must mean buf needs more data
  214. delta = self.socket.recv(8096)
  215. if delta == '':
  216. return None
  217. self._buf += delta
  218. msgs = self._parse_messages()
  219. self._msgs.extend(msgs)
  220. return self._msgs.popleft()
  221. def _send_closing_frame(self, ignore_send_errors=False):
  222. """Sends the closing frame to the client, if required."""
  223. if self.version == 76 and not self.websocket_closed:
  224. try:
  225. self.socket.sendall("\xff\x00")
  226. except SocketError:
  227. # Sometimes, like when the remote side cuts off the connection,
  228. # we don't care about this.
  229. if not ignore_send_errors: #pragma NO COVER
  230. raise
  231. self.websocket_closed = True
  232. def close(self):
  233. """Forcibly close the websocket; generally it is preferable to
  234. return from the handler method."""
  235. self._send_closing_frame()
  236. self.socket.shutdown(True)
  237. self.socket.close()