michael@0: # Copyright 2012, Google Inc.
michael@0: # All rights reserved.
michael@0: #
michael@0: # Redistribution and use in source and binary forms, with or without
michael@0: # modification, are permitted provided that the following conditions are
michael@0: # met:
michael@0: #
michael@0: #     * Redistributions of source code must retain the above copyright
michael@0: # notice, this list of conditions and the following disclaimer.
michael@0: #     * Redistributions in binary form must reproduce the above
michael@0: # copyright notice, this list of conditions and the following disclaimer
michael@0: # in the documentation and/or other materials provided with the
michael@0: # distribution.
michael@0: #     * Neither the name of Google Inc. nor the names of its
michael@0: # contributors may be used to endorse or promote products derived from
michael@0: # this software without specific prior written permission.
michael@0: #
michael@0: # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
michael@0: # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
michael@0: # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
michael@0: # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
michael@0: # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
michael@0: # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
michael@0: # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
michael@0: # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
michael@0: # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
michael@0: # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
michael@0: # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
michael@0: 
michael@0: 
michael@0: """Dispatch WebSocket request.
michael@0: """
michael@0: 
michael@0: 
michael@0: import logging
michael@0: import os
michael@0: import re
michael@0: 
michael@0: from mod_pywebsocket import common
michael@0: from mod_pywebsocket import handshake
michael@0: from mod_pywebsocket import msgutil
michael@0: from mod_pywebsocket import stream
michael@0: from mod_pywebsocket import util
michael@0: 
michael@0: 
michael@0: _SOURCE_PATH_PATTERN = re.compile(r'(?i)_wsh\.py$')
michael@0: _SOURCE_SUFFIX = '_wsh.py'
michael@0: _DO_EXTRA_HANDSHAKE_HANDLER_NAME = 'web_socket_do_extra_handshake'
michael@0: _TRANSFER_DATA_HANDLER_NAME = 'web_socket_transfer_data'
michael@0: _PASSIVE_CLOSING_HANDSHAKE_HANDLER_NAME = (
michael@0:     'web_socket_passive_closing_handshake')
michael@0: 
michael@0: 
michael@0: class DispatchException(Exception):
michael@0:     """Exception in dispatching WebSocket request."""
michael@0: 
michael@0:     def __init__(self, name, status=common.HTTP_STATUS_NOT_FOUND):
michael@0:         super(DispatchException, self).__init__(name)
michael@0:         self.status = status
michael@0: 
michael@0: 
michael@0: def _default_passive_closing_handshake_handler(request):
michael@0:     """Default web_socket_passive_closing_handshake handler."""
michael@0: 
michael@0:     return common.STATUS_NORMAL_CLOSURE, ''
michael@0: 
michael@0: 
michael@0: def _normalize_path(path):
michael@0:     """Normalize path.
michael@0: 
michael@0:     Args:
michael@0:         path: the path to normalize.
michael@0: 
michael@0:     Path is converted to the absolute path.
michael@0:     The input path can use either '\\' or '/' as the separator.
michael@0:     The normalized path always uses '/' regardless of the platform.
michael@0:     """
michael@0: 
michael@0:     path = path.replace('\\', os.path.sep)
michael@0:     path = os.path.realpath(path)
michael@0:     path = path.replace('\\', '/')
michael@0:     return path
michael@0: 
michael@0: 
michael@0: def _create_path_to_resource_converter(base_dir):
michael@0:     """Returns a function that converts the path of a WebSocket handler source
michael@0:     file to a resource string by removing the path to the base directory from
michael@0:     its head, removing _SOURCE_SUFFIX from its tail, and replacing path
michael@0:     separators in it with '/'.
michael@0: 
michael@0:     Args:
michael@0:         base_dir: the path to the base directory.
michael@0:     """
michael@0: 
michael@0:     base_dir = _normalize_path(base_dir)
michael@0: 
michael@0:     base_len = len(base_dir)
michael@0:     suffix_len = len(_SOURCE_SUFFIX)
michael@0: 
michael@0:     def converter(path):
michael@0:         if not path.endswith(_SOURCE_SUFFIX):
michael@0:             return None
michael@0:         # _normalize_path must not be used because resolving symlink breaks
michael@0:         # following path check.
michael@0:         path = path.replace('\\', '/')
michael@0:         if not path.startswith(base_dir):
michael@0:             return None
michael@0:         return path[base_len:-suffix_len]
michael@0: 
michael@0:     return converter
michael@0: 
michael@0: 
michael@0: def _enumerate_handler_file_paths(directory):
michael@0:     """Returns a generator that enumerates WebSocket Handler source file names
michael@0:     in the given directory.
michael@0:     """
michael@0: 
michael@0:     for root, unused_dirs, files in os.walk(directory):
michael@0:         for base in files:
michael@0:             path = os.path.join(root, base)
michael@0:             if _SOURCE_PATH_PATTERN.search(path):
michael@0:                 yield path
michael@0: 
michael@0: 
michael@0: class _HandlerSuite(object):
michael@0:     """A handler suite holder class."""
michael@0: 
michael@0:     def __init__(self, do_extra_handshake, transfer_data,
michael@0:                  passive_closing_handshake):
michael@0:         self.do_extra_handshake = do_extra_handshake
michael@0:         self.transfer_data = transfer_data
michael@0:         self.passive_closing_handshake = passive_closing_handshake
michael@0: 
michael@0: 
michael@0: def _source_handler_file(handler_definition):
michael@0:     """Source a handler definition string.
michael@0: 
michael@0:     Args:
michael@0:         handler_definition: a string containing Python statements that define
michael@0:                             handler functions.
michael@0:     """
michael@0: 
michael@0:     global_dic = {}
michael@0:     try:
michael@0:         exec handler_definition in global_dic
michael@0:     except Exception:
michael@0:         raise DispatchException('Error in sourcing handler:' +
michael@0:                                 util.get_stack_trace())
michael@0:     passive_closing_handshake_handler = None
michael@0:     try:
michael@0:         passive_closing_handshake_handler = _extract_handler(
michael@0:             global_dic, _PASSIVE_CLOSING_HANDSHAKE_HANDLER_NAME)
michael@0:     except Exception:
michael@0:         passive_closing_handshake_handler = (
michael@0:             _default_passive_closing_handshake_handler)
michael@0:     return _HandlerSuite(
michael@0:         _extract_handler(global_dic, _DO_EXTRA_HANDSHAKE_HANDLER_NAME),
michael@0:         _extract_handler(global_dic, _TRANSFER_DATA_HANDLER_NAME),
michael@0:         passive_closing_handshake_handler)
michael@0: 
michael@0: 
michael@0: def _extract_handler(dic, name):
michael@0:     """Extracts a callable with the specified name from the given dictionary
michael@0:     dic.
michael@0:     """
michael@0: 
michael@0:     if name not in dic:
michael@0:         raise DispatchException('%s is not defined.' % name)
michael@0:     handler = dic[name]
michael@0:     if not callable(handler):
michael@0:         raise DispatchException('%s is not callable.' % name)
michael@0:     return handler
michael@0: 
michael@0: 
michael@0: class Dispatcher(object):
michael@0:     """Dispatches WebSocket requests.
michael@0: 
michael@0:     This class maintains a map from resource name to handlers.
michael@0:     """
michael@0: 
michael@0:     def __init__(
michael@0:         self, root_dir, scan_dir=None,
michael@0:         allow_handlers_outside_root_dir=True):
michael@0:         """Construct an instance.
michael@0: 
michael@0:         Args:
michael@0:             root_dir: The directory where handler definition files are
michael@0:                       placed.
michael@0:             scan_dir: The directory where handler definition files are
michael@0:                       searched. scan_dir must be a directory under root_dir,
michael@0:                       including root_dir itself.  If scan_dir is None,
michael@0:                       root_dir is used as scan_dir. scan_dir can be useful
michael@0:                       in saving scan time when root_dir contains many
michael@0:                       subdirectories.
michael@0:             allow_handlers_outside_root_dir: Scans handler files even if their
michael@0:                       canonical path is not under root_dir.
michael@0:         """
michael@0: 
michael@0:         self._logger = util.get_class_logger(self)
michael@0: 
michael@0:         self._handler_suite_map = {}
michael@0:         self._source_warnings = []
michael@0:         if scan_dir is None:
michael@0:             scan_dir = root_dir
michael@0:         if not os.path.realpath(scan_dir).startswith(
michael@0:                 os.path.realpath(root_dir)):
michael@0:             raise DispatchException('scan_dir:%s must be a directory under '
michael@0:                                     'root_dir:%s.' % (scan_dir, root_dir))
michael@0:         self._source_handler_files_in_dir(
michael@0:             root_dir, scan_dir, allow_handlers_outside_root_dir)
michael@0: 
michael@0:     def add_resource_path_alias(self,
michael@0:                                 alias_resource_path, existing_resource_path):
michael@0:         """Add resource path alias.
michael@0: 
michael@0:         Once added, request to alias_resource_path would be handled by
michael@0:         handler registered for existing_resource_path.
michael@0: 
michael@0:         Args:
michael@0:             alias_resource_path: alias resource path
michael@0:             existing_resource_path: existing resource path
michael@0:         """
michael@0:         try:
michael@0:             handler_suite = self._handler_suite_map[existing_resource_path]
michael@0:             self._handler_suite_map[alias_resource_path] = handler_suite
michael@0:         except KeyError:
michael@0:             raise DispatchException('No handler for: %r' %
michael@0:                                     existing_resource_path)
michael@0: 
michael@0:     def source_warnings(self):
michael@0:         """Return warnings in sourcing handlers."""
michael@0: 
michael@0:         return self._source_warnings
michael@0: 
michael@0:     def do_extra_handshake(self, request):
michael@0:         """Do extra checking in WebSocket handshake.
michael@0: 
michael@0:         Select a handler based on request.uri and call its
michael@0:         web_socket_do_extra_handshake function.
michael@0: 
michael@0:         Args:
michael@0:             request: mod_python request.
michael@0: 
michael@0:         Raises:
michael@0:             DispatchException: when handler was not found
michael@0:             AbortedByUserException: when user handler abort connection
michael@0:             HandshakeException: when opening handshake failed
michael@0:         """
michael@0: 
michael@0:         handler_suite = self.get_handler_suite(request.ws_resource)
michael@0:         if handler_suite is None:
michael@0:             raise DispatchException('No handler for: %r' % request.ws_resource)
michael@0:         do_extra_handshake_ = handler_suite.do_extra_handshake
michael@0:         try:
michael@0:             do_extra_handshake_(request)
michael@0:         except handshake.AbortedByUserException, e:
michael@0:             raise
michael@0:         except Exception, e:
michael@0:             util.prepend_message_to_exception(
michael@0:                     '%s raised exception for %s: ' % (
michael@0:                             _DO_EXTRA_HANDSHAKE_HANDLER_NAME,
michael@0:                             request.ws_resource),
michael@0:                     e)
michael@0:             raise handshake.HandshakeException(e, common.HTTP_STATUS_FORBIDDEN)
michael@0: 
michael@0:     def transfer_data(self, request):
michael@0:         """Let a handler transfer_data with a WebSocket client.
michael@0: 
michael@0:         Select a handler based on request.ws_resource and call its
michael@0:         web_socket_transfer_data function.
michael@0: 
michael@0:         Args:
michael@0:             request: mod_python request.
michael@0: 
michael@0:         Raises:
michael@0:             DispatchException: when handler was not found
michael@0:             AbortedByUserException: when user handler abort connection
michael@0:         """
michael@0: 
michael@0:         handler_suite = self.get_handler_suite(request.ws_resource)
michael@0:         if handler_suite is None:
michael@0:             raise DispatchException('No handler for: %r' % request.ws_resource)
michael@0:         transfer_data_ = handler_suite.transfer_data
michael@0:         # TODO(tyoshino): Terminate underlying TCP connection if possible.
michael@0:         try:
michael@0:             transfer_data_(request)
michael@0:             if not request.server_terminated:
michael@0:                 request.ws_stream.close_connection()
michael@0:         # Catch non-critical exceptions the handler didn't handle.
michael@0:         except handshake.AbortedByUserException, e:
michael@0:             self._logger.debug('%s', e)
michael@0:             raise
michael@0:         except msgutil.BadOperationException, e:
michael@0:             self._logger.debug('%s', e)
michael@0:             request.ws_stream.close_connection(common.STATUS_ABNORMAL_CLOSURE)
michael@0:         except msgutil.InvalidFrameException, e:
michael@0:             # InvalidFrameException must be caught before
michael@0:             # ConnectionTerminatedException that catches InvalidFrameException.
michael@0:             self._logger.debug('%s', e)
michael@0:             request.ws_stream.close_connection(common.STATUS_PROTOCOL_ERROR)
michael@0:         except msgutil.UnsupportedFrameException, e:
michael@0:             self._logger.debug('%s', e)
michael@0:             request.ws_stream.close_connection(common.STATUS_UNSUPPORTED_DATA)
michael@0:         except stream.InvalidUTF8Exception, e:
michael@0:             self._logger.debug('%s', e)
michael@0:             request.ws_stream.close_connection(
michael@0:                 common.STATUS_INVALID_FRAME_PAYLOAD_DATA)
michael@0:         except msgutil.ConnectionTerminatedException, e:
michael@0:             self._logger.debug('%s', e)
michael@0:         except Exception, e:
michael@0:             util.prepend_message_to_exception(
michael@0:                 '%s raised exception for %s: ' % (
michael@0:                     _TRANSFER_DATA_HANDLER_NAME, request.ws_resource),
michael@0:                 e)
michael@0:             raise
michael@0: 
michael@0:     def passive_closing_handshake(self, request):
michael@0:         """Prepare code and reason for responding client initiated closing
michael@0:         handshake.
michael@0:         """
michael@0: 
michael@0:         handler_suite = self.get_handler_suite(request.ws_resource)
michael@0:         if handler_suite is None:
michael@0:             return _default_passive_closing_handshake_handler(request)
michael@0:         return handler_suite.passive_closing_handshake(request)
michael@0: 
michael@0:     def get_handler_suite(self, resource):
michael@0:         """Retrieves two handlers (one for extra handshake processing, and one
michael@0:         for data transfer) for the given request as a HandlerSuite object.
michael@0:         """
michael@0: 
michael@0:         fragment = None
michael@0:         if '#' in resource:
michael@0:             resource, fragment = resource.split('#', 1)
michael@0:         if '?' in resource:
michael@0:             resource = resource.split('?', 1)[0]
michael@0:         handler_suite = self._handler_suite_map.get(resource)
michael@0:         if handler_suite and fragment:
michael@0:             raise DispatchException('Fragment identifiers MUST NOT be used on '
michael@0:                                     'WebSocket URIs',
michael@0:                                     common.HTTP_STATUS_BAD_REQUEST)
michael@0:         return handler_suite
michael@0: 
michael@0:     def _source_handler_files_in_dir(
michael@0:         self, root_dir, scan_dir, allow_handlers_outside_root_dir):
michael@0:         """Source all the handler source files in the scan_dir directory.
michael@0: 
michael@0:         The resource path is determined relative to root_dir.
michael@0:         """
michael@0: 
michael@0:         # We build a map from resource to handler code assuming that there's
michael@0:         # only one path from root_dir to scan_dir and it can be obtained by
michael@0:         # comparing realpath of them.
michael@0: 
michael@0:         # Here we cannot use abspath. See
michael@0:         # https://bugs.webkit.org/show_bug.cgi?id=31603
michael@0: 
michael@0:         convert = _create_path_to_resource_converter(root_dir)
michael@0:         scan_realpath = os.path.realpath(scan_dir)
michael@0:         root_realpath = os.path.realpath(root_dir)
michael@0:         for path in _enumerate_handler_file_paths(scan_realpath):
michael@0:             if (not allow_handlers_outside_root_dir and
michael@0:                 (not os.path.realpath(path).startswith(root_realpath))):
michael@0:                 self._logger.debug(
michael@0:                     'Canonical path of %s is not under root directory' %
michael@0:                     path)
michael@0:                 continue
michael@0:             try:
michael@0:                 handler_suite = _source_handler_file(open(path).read())
michael@0:             except DispatchException, e:
michael@0:                 self._source_warnings.append('%s: %s' % (path, e))
michael@0:                 continue
michael@0:             resource = convert(path)
michael@0:             if resource is None:
michael@0:                 self._logger.debug(
michael@0:                     'Path to resource conversion on %s failed' % path)
michael@0:             else:
michael@0:                 self._handler_suite_map[convert(path)] = handler_suite
michael@0: 
michael@0: 
michael@0: # vi:sts=4 sw=4 et