1# Copyright 2012, Google Inc.
2# All rights reserved.
3#
4# Redistribution and use in source and binary forms, with or without
5# modification, are permitted provided that the following conditions are
6# met:
7#
8#     * Redistributions of source code must retain the above copyright
9# notice, this list of conditions and the following disclaimer.
10#     * Redistributions in binary form must reproduce the above
11# copyright notice, this list of conditions and the following disclaimer
12# in the documentation and/or other materials provided with the
13# distribution.
14#     * Neither the name of Google Inc. nor the names of its
15# contributors may be used to endorse or promote products derived from
16# this software without specific prior written permission.
17#
18# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29
30
31"""Dispatch WebSocket request.
32"""
33
34
35import logging
36import os
37import re
38
39from mod_pywebsocket import common
40from mod_pywebsocket import handshake
41from mod_pywebsocket import msgutil
42from mod_pywebsocket import mux
43from mod_pywebsocket import stream
44from mod_pywebsocket import util
45
46
47_SOURCE_PATH_PATTERN = re.compile(r'(?i)_wsh\.py$')
48_SOURCE_SUFFIX = '_wsh.py'
49_DO_EXTRA_HANDSHAKE_HANDLER_NAME = 'web_socket_do_extra_handshake'
50_TRANSFER_DATA_HANDLER_NAME = 'web_socket_transfer_data'
51_PASSIVE_CLOSING_HANDSHAKE_HANDLER_NAME = (
52    'web_socket_passive_closing_handshake')
53
54
55class DispatchException(Exception):
56    """Exception in dispatching WebSocket request."""
57
58    def __init__(self, name, status=common.HTTP_STATUS_NOT_FOUND):
59        super(DispatchException, self).__init__(name)
60        self.status = status
61
62
63def _default_passive_closing_handshake_handler(request):
64    """Default web_socket_passive_closing_handshake handler."""
65
66    return common.STATUS_NORMAL_CLOSURE, ''
67
68
69def _normalize_path(path):
70    """Normalize path.
71
72    Args:
73        path: the path to normalize.
74
75    Path is converted to the absolute path.
76    The input path can use either '\\' or '/' as the separator.
77    The normalized path always uses '/' regardless of the platform.
78    """
79
80    path = path.replace('\\', os.path.sep)
81    path = os.path.realpath(path)
82    path = path.replace('\\', '/')
83    return path
84
85
86def _create_path_to_resource_converter(base_dir):
87    """Returns a function that converts the path of a WebSocket handler source
88    file to a resource string by removing the path to the base directory from
89    its head, removing _SOURCE_SUFFIX from its tail, and replacing path
90    separators in it with '/'.
91
92    Args:
93        base_dir: the path to the base directory.
94    """
95
96    base_dir = _normalize_path(base_dir)
97
98    base_len = len(base_dir)
99    suffix_len = len(_SOURCE_SUFFIX)
100
101    def converter(path):
102        if not path.endswith(_SOURCE_SUFFIX):
103            return None
104        # _normalize_path must not be used because resolving symlink breaks
105        # following path check.
106        path = path.replace('\\', '/')
107        if not path.startswith(base_dir):
108            return None
109        return path[base_len:-suffix_len]
110
111    return converter
112
113
114def _enumerate_handler_file_paths(directory):
115    """Returns a generator that enumerates WebSocket Handler source file names
116    in the given directory.
117    """
118
119    for root, unused_dirs, files in os.walk(directory):
120        for base in files:
121            path = os.path.join(root, base)
122            if _SOURCE_PATH_PATTERN.search(path):
123                yield path
124
125
126class _HandlerSuite(object):
127    """A handler suite holder class."""
128
129    def __init__(self, do_extra_handshake, transfer_data,
130                 passive_closing_handshake):
131        self.do_extra_handshake = do_extra_handshake
132        self.transfer_data = transfer_data
133        self.passive_closing_handshake = passive_closing_handshake
134
135
136def _source_handler_file(handler_definition):
137    """Source a handler definition string.
138
139    Args:
140        handler_definition: a string containing Python statements that define
141                            handler functions.
142    """
143
144    global_dic = {}
145    try:
146        exec handler_definition in global_dic
147    except Exception:
148        raise DispatchException('Error in sourcing handler:' +
149                                util.get_stack_trace())
150    passive_closing_handshake_handler = None
151    try:
152        passive_closing_handshake_handler = _extract_handler(
153            global_dic, _PASSIVE_CLOSING_HANDSHAKE_HANDLER_NAME)
154    except Exception:
155        passive_closing_handshake_handler = (
156            _default_passive_closing_handshake_handler)
157    return _HandlerSuite(
158        _extract_handler(global_dic, _DO_EXTRA_HANDSHAKE_HANDLER_NAME),
159        _extract_handler(global_dic, _TRANSFER_DATA_HANDLER_NAME),
160        passive_closing_handshake_handler)
161
162
163def _extract_handler(dic, name):
164    """Extracts a callable with the specified name from the given dictionary
165    dic.
166    """
167
168    if name not in dic:
169        raise DispatchException('%s is not defined.' % name)
170    handler = dic[name]
171    if not callable(handler):
172        raise DispatchException('%s is not callable.' % name)
173    return handler
174
175
176class Dispatcher(object):
177    """Dispatches WebSocket requests.
178
179    This class maintains a map from resource name to handlers.
180    """
181
182    def __init__(
183        self, root_dir, scan_dir=None,
184        allow_handlers_outside_root_dir=True):
185        """Construct an instance.
186
187        Args:
188            root_dir: The directory where handler definition files are
189                      placed.
190            scan_dir: The directory where handler definition files are
191                      searched. scan_dir must be a directory under root_dir,
192                      including root_dir itself.  If scan_dir is None,
193                      root_dir is used as scan_dir. scan_dir can be useful
194                      in saving scan time when root_dir contains many
195                      subdirectories.
196            allow_handlers_outside_root_dir: Scans handler files even if their
197                      canonical path is not under root_dir.
198        """
199
200        self._logger = util.get_class_logger(self)
201
202        self._handler_suite_map = {}
203        self._source_warnings = []
204        if scan_dir is None:
205            scan_dir = root_dir
206        if not os.path.realpath(scan_dir).startswith(
207                os.path.realpath(root_dir)):
208            raise DispatchException('scan_dir:%s must be a directory under '
209                                    'root_dir:%s.' % (scan_dir, root_dir))
210        self._source_handler_files_in_dir(
211            root_dir, scan_dir, allow_handlers_outside_root_dir)
212
213    def add_resource_path_alias(self,
214                                alias_resource_path, existing_resource_path):
215        """Add resource path alias.
216
217        Once added, request to alias_resource_path would be handled by
218        handler registered for existing_resource_path.
219
220        Args:
221            alias_resource_path: alias resource path
222            existing_resource_path: existing resource path
223        """
224        try:
225            handler_suite = self._handler_suite_map[existing_resource_path]
226            self._handler_suite_map[alias_resource_path] = handler_suite
227        except KeyError:
228            raise DispatchException('No handler for: %r' %
229                                    existing_resource_path)
230
231    def source_warnings(self):
232        """Return warnings in sourcing handlers."""
233
234        return self._source_warnings
235
236    def do_extra_handshake(self, request):
237        """Do extra checking in WebSocket handshake.
238
239        Select a handler based on request.uri and call its
240        web_socket_do_extra_handshake function.
241
242        Args:
243            request: mod_python request.
244
245        Raises:
246            DispatchException: when handler was not found
247            AbortedByUserException: when user handler abort connection
248            HandshakeException: when opening handshake failed
249        """
250
251        handler_suite = self.get_handler_suite(request.ws_resource)
252        if handler_suite is None:
253            raise DispatchException('No handler for: %r' % request.ws_resource)
254        do_extra_handshake_ = handler_suite.do_extra_handshake
255        try:
256            do_extra_handshake_(request)
257        except handshake.AbortedByUserException, e:
258            # Re-raise to tell the caller of this function to finish this
259            # connection without sending any error.
260            self._logger.debug('%s', util.get_stack_trace())
261            raise
262        except Exception, e:
263            util.prepend_message_to_exception(
264                    '%s raised exception for %s: ' % (
265                            _DO_EXTRA_HANDSHAKE_HANDLER_NAME,
266                            request.ws_resource),
267                    e)
268            raise handshake.HandshakeException(e, common.HTTP_STATUS_FORBIDDEN)
269
270    def transfer_data(self, request):
271        """Let a handler transfer_data with a WebSocket client.
272
273        Select a handler based on request.ws_resource and call its
274        web_socket_transfer_data function.
275
276        Args:
277            request: mod_python request.
278
279        Raises:
280            DispatchException: when handler was not found
281            AbortedByUserException: when user handler abort connection
282        """
283
284        # TODO(tyoshino): Terminate underlying TCP connection if possible.
285        try:
286            if mux.use_mux(request):
287                mux.start(request, self)
288            else:
289                handler_suite = self.get_handler_suite(request.ws_resource)
290                if handler_suite is None:
291                    raise DispatchException('No handler for: %r' %
292                                            request.ws_resource)
293                transfer_data_ = handler_suite.transfer_data
294                transfer_data_(request)
295
296            if not request.server_terminated:
297                request.ws_stream.close_connection()
298        # Catch non-critical exceptions the handler didn't handle.
299        except handshake.AbortedByUserException, e:
300            self._logger.debug('%s', util.get_stack_trace())
301            raise
302        except msgutil.BadOperationException, e:
303            self._logger.debug('%s', e)
304            request.ws_stream.close_connection(
305                common.STATUS_INTERNAL_ENDPOINT_ERROR)
306        except msgutil.InvalidFrameException, e:
307            # InvalidFrameException must be caught before
308            # ConnectionTerminatedException that catches InvalidFrameException.
309            self._logger.debug('%s', e)
310            request.ws_stream.close_connection(common.STATUS_PROTOCOL_ERROR)
311        except msgutil.UnsupportedFrameException, e:
312            self._logger.debug('%s', e)
313            request.ws_stream.close_connection(common.STATUS_UNSUPPORTED_DATA)
314        except stream.InvalidUTF8Exception, e:
315            self._logger.debug('%s', e)
316            request.ws_stream.close_connection(
317                common.STATUS_INVALID_FRAME_PAYLOAD_DATA)
318        except msgutil.ConnectionTerminatedException, e:
319            self._logger.debug('%s', e)
320        except Exception, e:
321            # Any other exceptions are forwarded to the caller of this
322            # function.
323            util.prepend_message_to_exception(
324                '%s raised exception for %s: ' % (
325                    _TRANSFER_DATA_HANDLER_NAME, request.ws_resource),
326                e)
327            raise
328
329    def passive_closing_handshake(self, request):
330        """Prepare code and reason for responding client initiated closing
331        handshake.
332        """
333
334        handler_suite = self.get_handler_suite(request.ws_resource)
335        if handler_suite is None:
336            return _default_passive_closing_handshake_handler(request)
337        return handler_suite.passive_closing_handshake(request)
338
339    def get_handler_suite(self, resource):
340        """Retrieves two handlers (one for extra handshake processing, and one
341        for data transfer) for the given request as a HandlerSuite object.
342        """
343
344        fragment = None
345        if '#' in resource:
346            resource, fragment = resource.split('#', 1)
347        if '?' in resource:
348            resource = resource.split('?', 1)[0]
349        handler_suite = self._handler_suite_map.get(resource)
350        if handler_suite and fragment:
351            raise DispatchException('Fragment identifiers MUST NOT be used on '
352                                    'WebSocket URIs',
353                                    common.HTTP_STATUS_BAD_REQUEST)
354        return handler_suite
355
356    def _source_handler_files_in_dir(
357        self, root_dir, scan_dir, allow_handlers_outside_root_dir):
358        """Source all the handler source files in the scan_dir directory.
359
360        The resource path is determined relative to root_dir.
361        """
362
363        # We build a map from resource to handler code assuming that there's
364        # only one path from root_dir to scan_dir and it can be obtained by
365        # comparing realpath of them.
366
367        # Here we cannot use abspath. See
368        # https://bugs.webkit.org/show_bug.cgi?id=31603
369
370        convert = _create_path_to_resource_converter(root_dir)
371        scan_realpath = os.path.realpath(scan_dir)
372        root_realpath = os.path.realpath(root_dir)
373        for path in _enumerate_handler_file_paths(scan_realpath):
374            if (not allow_handlers_outside_root_dir and
375                (not os.path.realpath(path).startswith(root_realpath))):
376                self._logger.debug(
377                    'Canonical path of %s is not under root directory' %
378                    path)
379                continue
380            try:
381                handler_suite = _source_handler_file(open(path).read())
382            except DispatchException, e:
383                self._source_warnings.append('%s: %s' % (path, e))
384                continue
385            resource = convert(path)
386            if resource is None:
387                self._logger.debug(
388                    'Path to resource conversion on %s failed' % path)
389            else:
390                self._handler_suite_map[convert(path)] = handler_suite
391
392
393# vi:sts=4 sw=4 et
394