PageRenderTime 50ms CodeModel.GetById 17ms RepoModel.GetById 1ms app.codeStats 0ms

/tornadio2/persistent.py

http://github.com/mrjoes/tornadio2
Python | 193 lines | 165 code | 7 blank | 21 comment | 0 complexity | 2c8d1b204fea5ae47ac338b52d801d4e MD5 | raw file
Possible License(s): Apache-2.0
  1. # -*- coding: utf-8 -*-
  2. #
  3. # Copyright: (c) 2011 by the Serge S. Koval, see AUTHORS for more details.
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License"); you may
  6. # not use this file except in compliance with the License. You may obtain
  7. # a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  13. # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  14. # License for the specific language governing permissions and limitations
  15. # under the License.
  16. """
  17. tornadio2.persistent
  18. ~~~~~~~~~~~~~~~~~~~~
  19. Persistent transport implementations.
  20. """
  21. import logging
  22. import time
  23. import traceback
  24. import tornado
  25. from tornado.web import HTTPError
  26. from tornado import stack_context
  27. from tornado.websocket import WebSocketHandler
  28. from tornadio2 import proto
  29. logger = logging.getLogger('tornadio2.persistent')
  30. class TornadioWebSocketHandler(WebSocketHandler):
  31. """Websocket protocol handler"""
  32. # Transport name
  33. name = 'websocket'
  34. def initialize(self, server):
  35. self.server = server
  36. self.session = None
  37. self._is_active = not self.server.settings['websocket_check']
  38. self._global_heartbeats = self.server.settings['global_heartbeats']
  39. logger.debug('Initializing %s handler.' % self.name)
  40. # Additional verification of the websocket handshake
  41. # For now it will stay here, till https://github.com/facebook/tornado/pull/415
  42. # is merged.
  43. def _execute(self, transforms, *args, **kwargs):
  44. with stack_context.ExceptionStackContext(self._handle_websocket_exception):
  45. # Websocket only supports GET method
  46. if self.request.method != 'GET':
  47. self.stream.write(tornado.escape.utf8(
  48. "HTTP/1.1 405 Method Not Allowed\r\n\r\n"
  49. ))
  50. self.stream.close()
  51. return
  52. # Upgrade header should be present and should be equal to WebSocket
  53. if self.request.headers.get("Upgrade", "").lower() != 'websocket':
  54. self.stream.write(tornado.escape.utf8(
  55. "HTTP/1.1 400 Bad Request\r\n\r\n"
  56. "Can \"Upgrade\" only to \"WebSocket\"."
  57. ))
  58. self.stream.close()
  59. return
  60. # Connection header should be upgrade. Some proxy servers/load balancers
  61. # might mess with it.
  62. if self.request.headers.get("Connection", "").lower().find('upgrade') == -1:
  63. self.stream.write(tornado.escape.utf8(
  64. "HTTP/1.1 400 Bad Request\r\n\r\n"
  65. "\"Connection\" must be \"Upgrade\"."
  66. ))
  67. self.stream.close()
  68. return
  69. super(TornadioWebSocketHandler, self)._execute(transforms, *args, **kwargs)
  70. def open(self, session_id):
  71. """WebSocket open handler"""
  72. self.session = self.server.get_session(session_id)
  73. if self.session is None:
  74. raise HTTPError(401, "Invalid Session")
  75. if not self._is_active:
  76. # Need to check if websocket connection was really established by sending hearbeat packet
  77. # and waiting for response
  78. self.write_message(proto.heartbeat())
  79. self.server.io_loop.add_timeout(time.time() + self.server.settings['client_timeout'],
  80. self._connection_check)
  81. else:
  82. # Associate session handler
  83. self.session.set_handler(self)
  84. self.session.reset_heartbeat()
  85. # Flush messages, if any
  86. self.session.flush()
  87. def _connection_check(self):
  88. if not self._is_active:
  89. self._detach()
  90. try:
  91. # Might throw exception if connection was closed already
  92. self.close()
  93. except:
  94. pass
  95. def _detach(self):
  96. if self.session is not None:
  97. if self._is_active:
  98. self.session.stop_heartbeat()
  99. self.session.remove_handler(self)
  100. self.session = None
  101. def on_message(self, message):
  102. # Tracking
  103. self.server.stats.on_packet_recv(1)
  104. # Fix for late messages (after connection was closed)
  105. if not self.session:
  106. return
  107. # Mark that connection is active and flush any pending messages
  108. if not self._is_active:
  109. # Associate session handler and flush queued messages
  110. self.session.set_handler(self)
  111. self.session.reset_heartbeat()
  112. self.session.flush()
  113. self._is_active = True
  114. if not self._global_heartbeats:
  115. self.session.delay_heartbeat()
  116. try:
  117. self.session.raw_message(message)
  118. except Exception, ex:
  119. logger.error('Failed to handle message: ' + traceback.format_exc(ex))
  120. # Close session on exception
  121. if self.session is not None:
  122. self.session.close()
  123. def on_close(self):
  124. self._detach()
  125. def send_messages(self, messages):
  126. # Tracking
  127. self.server.stats.on_packet_sent(len(messages))
  128. try:
  129. for m in messages:
  130. self.write_message(m)
  131. except IOError:
  132. if self.ws_connection and self.ws_connection.client_terminated:
  133. logger.debug('Dropping active websocket connection due to IOError.')
  134. self._detach()
  135. def session_closed(self):
  136. try:
  137. self.close()
  138. except Exception:
  139. logger.debug('Exception', exc_info=True)
  140. finally:
  141. self._detach()
  142. def _handle_websocket_exception(self, type, value, traceback):
  143. if type is IOError:
  144. self.server.io_loop.add_callback(self.on_connection_close)
  145. # raise (type, value, traceback)
  146. logger.debug('Exception', exc_info=(type, value, traceback))
  147. return True
  148. # Websocket overrides
  149. def allow_draft76(self):
  150. return True
  151. class TornadioFlashSocketHandler(TornadioWebSocketHandler):
  152. # Transport name
  153. name = 'flashsocket'