Home | History | Annotate | Download | only in test 
      1 #!/usr/bin/env python
      2 #
      3 # Copyright 2011, Google Inc.
      4 # All rights reserved.
      5 #
      6 # Redistribution and use in source and binary forms, with or without
      7 # modification, are permitted provided that the following conditions are
      8 # met:
      9 #
     10 #     * Redistributions of source code must retain the above copyright
     11 # notice, this list of conditions and the following disclaimer.
     12 #     * Redistributions in binary form must reproduce the above
     13 # copyright notice, this list of conditions and the following disclaimer
     14 # in the documentation and/or other materials provided with the
     15 # distribution.
     16 #     * Neither the name of Google Inc. nor the names of its
     17 # contributors may be used to endorse or promote products derived from
     18 # this software without specific prior written permission.
     19 #
     20 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
     21 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
     22 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
     23 # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
     24 # OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
     25 # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
     26 # LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
     27 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
     28 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
     29 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
     30 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
     31 
     32 
     33 """Tests for handshake module."""
     34 
     35 
     36 import unittest
     37 
     38 import set_sys_path  # Update sys.path to locate mod_pywebsocket module.
     39 from mod_pywebsocket import common
     40 from mod_pywebsocket.handshake._base import AbortedByUserException
     41 from mod_pywebsocket.handshake._base import HandshakeException
     42 from mod_pywebsocket.handshake._base import VersionException
     43 from mod_pywebsocket.handshake.hybi import Handshaker
     44 
     45 import mock
     46 
     47 
     48 class RequestDefinition(object):
     49     """A class for holding data for constructing opening handshake strings for
     50     testing the opening handshake processor.
     51     """
     52 
     53     def __init__(self, method, uri, headers):
     54         self.method = method
     55         self.uri = uri
     56         self.headers = headers
     57 
     58 
     59 def _create_good_request_def():
     60     return RequestDefinition(
     61         'GET', '/demo',
     62         {'Host': 'server.example.com',
     63          'Upgrade': 'websocket',
     64          'Connection': 'Upgrade',
     65          'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ==',
     66          'Sec-WebSocket-Origin': 'http://example.com',
     67          'Sec-WebSocket-Version': '8'})
     68 
     69 
     70 def _create_request(request_def):
     71     conn = mock.MockConn('')
     72     return mock.MockRequest(
     73         method=request_def.method,
     74         uri=request_def.uri,
     75         headers_in=request_def.headers,
     76         connection=conn)
     77 
     78 
     79 def _create_handshaker(request):
     80     handshaker = Handshaker(request, mock.MockDispatcher())
     81     return handshaker
     82 
     83 
     84 class SubprotocolChoosingDispatcher(object):
     85     """A dispatcher for testing. This dispatcher sets the i-th subprotocol
     86     of requested ones to ws_protocol where i is given on construction as index
     87     argument. If index is negative, default_value will be set to ws_protocol.
     88     """
     89 
     90     def __init__(self, index, default_value=None):
     91         self.index = index
     92         self.default_value = default_value
     93 
     94     def do_extra_handshake(self, conn_context):
     95         if self.index >= 0:
     96             conn_context.ws_protocol = conn_context.ws_requested_protocols[
     97                 self.index]
     98         else:
     99             conn_context.ws_protocol = self.default_value
    100 
    101     def transfer_data(self, conn_context):
    102         pass
    103 
    104 
    105 class HandshakeAbortedException(Exception):
    106     pass
    107 
    108 
    109 class AbortingDispatcher(object):
    110     """A dispatcher for testing. This dispatcher raises an exception in
    111     do_extra_handshake to reject the request.
    112     """
    113 
    114     def do_extra_handshake(self, conn_context):
    115         raise HandshakeAbortedException('An exception to reject the request')
    116 
    117     def transfer_data(self, conn_context):
    118         pass
    119 
    120 
    121 class AbortedByUserDispatcher(object):
    122     """A dispatcher for testing. This dispatcher raises an
    123     AbortedByUserException in do_extra_handshake to reject the request.
    124     """
    125 
    126     def do_extra_handshake(self, conn_context):
    127         raise AbortedByUserException('An AbortedByUserException to reject the '
    128                                      'request')
    129 
    130     def transfer_data(self, conn_context):
    131         pass
    132 
    133 
    134 _EXPECTED_RESPONSE = (
    135     'HTTP/1.1 101 Switching Protocols\r\n'
    136     'Upgrade: websocket\r\n'
    137     'Connection: Upgrade\r\n'
    138     'Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n')
    139 
    140 
    141 class HandshakerTest(unittest.TestCase):
    142     """A unittest for draft-ietf-hybi-thewebsocketprotocol-06 and later
    143     handshake processor.
    144     """
    145 
    146     def test_do_handshake(self):
    147         request = _create_request(_create_good_request_def())
    148         dispatcher = mock.MockDispatcher()
    149         handshaker = Handshaker(request, dispatcher)
    150         handshaker.do_handshake()
    151 
    152         self.assertTrue(dispatcher.do_extra_handshake_called)
    153 
    154         self.assertEqual(
    155             _EXPECTED_RESPONSE, request.connection.written_data())
    156         self.assertEqual('/demo', request.ws_resource)
    157         self.assertEqual('http://example.com', request.ws_origin)
    158         self.assertEqual(None, request.ws_protocol)
    159         self.assertEqual(None, request.ws_extensions)
    160         self.assertEqual(common.VERSION_HYBI08, request.ws_version)
    161 
    162     def test_do_handshake_with_capitalized_value(self):
    163         request_def = _create_good_request_def()
    164         request_def.headers['upgrade'] = 'WEBSOCKET'
    165 
    166         request = _create_request(request_def)
    167         handshaker = _create_handshaker(request)
    168         handshaker.do_handshake()
    169         self.assertEqual(
    170             _EXPECTED_RESPONSE, request.connection.written_data())
    171 
    172         request_def = _create_good_request_def()
    173         request_def.headers['Connection'] = 'UPGRADE'
    174 
    175         request = _create_request(request_def)
    176         handshaker = _create_handshaker(request)
    177         handshaker.do_handshake()
    178         self.assertEqual(
    179             _EXPECTED_RESPONSE, request.connection.written_data())
    180 
    181     def test_do_handshake_with_multiple_connection_values(self):
    182         request_def = _create_good_request_def()
    183         request_def.headers['Connection'] = 'Upgrade, keep-alive, , '
    184 
    185         request = _create_request(request_def)
    186         handshaker = _create_handshaker(request)
    187         handshaker.do_handshake()
    188         self.assertEqual(
    189             _EXPECTED_RESPONSE, request.connection.written_data())
    190 
    191     def test_aborting_handshake(self):
    192         handshaker = Handshaker(
    193             _create_request(_create_good_request_def()),
    194             AbortingDispatcher())
    195         # do_extra_handshake raises an exception. Check that it's not caught by
    196         # do_handshake.
    197         self.assertRaises(HandshakeAbortedException, handshaker.do_handshake)
    198 
    199     def test_do_handshake_with_protocol(self):
    200         request_def = _create_good_request_def()
    201         request_def.headers['Sec-WebSocket-Protocol'] = 'chat, superchat'
    202 
    203         request = _create_request(request_def)
    204         handshaker = Handshaker(request, SubprotocolChoosingDispatcher(0))
    205         handshaker.do_handshake()
    206 
    207         EXPECTED_RESPONSE = (
    208             'HTTP/1.1 101 Switching Protocols\r\n'
    209             'Upgrade: websocket\r\n'
    210             'Connection: Upgrade\r\n'
    211             'Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n'
    212             'Sec-WebSocket-Protocol: chat\r\n\r\n')
    213 
    214         self.assertEqual(EXPECTED_RESPONSE, request.connection.written_data())
    215         self.assertEqual('chat', request.ws_protocol)
    216 
    217     def test_do_handshake_protocol_not_in_request_but_in_response(self):
    218         request_def = _create_good_request_def()
    219         request = _create_request(request_def)
    220         handshaker = Handshaker(
    221             request, SubprotocolChoosingDispatcher(-1, 'foobar'))
    222         # No request has been made but ws_protocol is set. HandshakeException
    223         # must be raised.
    224         self.assertRaises(HandshakeException, handshaker.do_handshake)
    225 
    226     def test_do_handshake_with_protocol_no_protocol_selection(self):
    227         request_def = _create_good_request_def()
    228         request_def.headers['Sec-WebSocket-Protocol'] = 'chat, superchat'
    229 
    230         request = _create_request(request_def)
    231         handshaker = _create_handshaker(request)
    232         # ws_protocol is not set. HandshakeException must be raised.
    233         self.assertRaises(HandshakeException, handshaker.do_handshake)
    234 
    235     def test_do_handshake_with_extensions(self):
    236         request_def = _create_good_request_def()
    237         request_def.headers['Sec-WebSocket-Extensions'] = (
    238             'deflate-stream, unknown')
    239 
    240         EXPECTED_RESPONSE = (
    241             'HTTP/1.1 101 Switching Protocols\r\n'
    242             'Upgrade: websocket\r\n'
    243             'Connection: Upgrade\r\n'
    244             'Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n'
    245             'Sec-WebSocket-Extensions: deflate-stream\r\n\r\n')
    246 
    247         request = _create_request(request_def)
    248         handshaker = _create_handshaker(request)
    249         handshaker.do_handshake()
    250         self.assertEqual(EXPECTED_RESPONSE, request.connection.written_data())
    251         self.assertEqual(1, len(request.ws_extensions))
    252         extension = request.ws_extensions[0]
    253         self.assertEqual('deflate-stream', extension.name())
    254         self.assertEqual(0, len(extension.get_parameter_names()))
    255 
    256     def test_do_handshake_with_quoted_extensions(self):
    257         request_def = _create_good_request_def()
    258         request_def.headers['Sec-WebSocket-Extensions'] = (
    259             'deflate-stream, , '
    260             'unknown; e   =    "mc^2"; ma="\r\n      \\\rf  "; pv=nrt')
    261 
    262         request = _create_request(request_def)
    263         handshaker = _create_handshaker(request)
    264         self.assertRaises(HandshakeException, handshaker.do_handshake)
    265 
    266     def test_do_handshake_with_optional_headers(self):
    267         request_def = _create_good_request_def()
    268         request_def.headers['EmptyValue'] = ''
    269         request_def.headers['AKey'] = 'AValue'
    270 
    271         request = _create_request(request_def)
    272         handshaker = _create_handshaker(request)
    273         handshaker.do_handshake()
    274         self.assertEqual(
    275             'AValue', request.headers_in['AKey'])
    276         self.assertEqual(
    277             '', request.headers_in['EmptyValue'])
    278 
    279     def test_abort_extra_handshake(self):
    280         handshaker = Handshaker(
    281             _create_request(_create_good_request_def()),
    282             AbortedByUserDispatcher())
    283         # do_extra_handshake raises an AbortedByUserException. Check that it's
    284         # not caught by do_handshake.
    285         self.assertRaises(AbortedByUserException, handshaker.do_handshake)
    286 
    287     def test_bad_requests(self):
    288         bad_cases = [
    289             ('HTTP request',
    290              RequestDefinition(
    291                  'GET', '/demo',
    292                  {'Host': 'www.google.com',
    293                   'User-Agent':
    294                       'Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.5;'
    295                       ' en-US; rv:1.9.1.3) Gecko/20090824 Firefox/3.5.3'
    296                       ' GTB6 GTBA',
    297                   'Accept':
    298                       'text/html,application/xhtml+xml,application/xml;q=0.9,'
    299                       '*/*;q=0.8',
    300                   'Accept-Language': 'en-us,en;q=0.5',
    301                   'Accept-Encoding': 'gzip,deflate',
    302                   'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
    303                   'Keep-Alive': '300',
    304                   'Connection': 'keep-alive'}), None, True)]
    305 
    306         request_def = _create_good_request_def()
    307         request_def.method = 'POST'
    308         bad_cases.append(('Wrong method', request_def, None, True))
    309 
    310         request_def = _create_good_request_def()
    311         del request_def.headers['Host']
    312         bad_cases.append(('Missing Host', request_def, None, True))
    313 
    314         request_def = _create_good_request_def()
    315         del request_def.headers['Upgrade']
    316         bad_cases.append(('Missing Upgrade', request_def, None, True))
    317 
    318         request_def = _create_good_request_def()
    319         request_def.headers['Upgrade'] = 'nonwebsocket'
    320         bad_cases.append(('Wrong Upgrade', request_def, None, True))
    321 
    322         request_def = _create_good_request_def()
    323         del request_def.headers['Connection']
    324         bad_cases.append(('Missing Connection', request_def, None, True))
    325 
    326         request_def = _create_good_request_def()
    327         request_def.headers['Connection'] = 'Downgrade'
    328         bad_cases.append(('Wrong Connection', request_def, None, True))
    329 
    330         request_def = _create_good_request_def()
    331         del request_def.headers['Sec-WebSocket-Key']
    332         bad_cases.append(('Missing Sec-WebSocket-Key', request_def, 400, True))
    333 
    334         request_def = _create_good_request_def()
    335         request_def.headers['Sec-WebSocket-Key'] = (
    336             'dGhlIHNhbXBsZSBub25jZQ==garbage')
    337         bad_cases.append(('Wrong Sec-WebSocket-Key (with garbage on the tail)',
    338                           request_def, 400, True))
    339 
    340         request_def = _create_good_request_def()
    341         request_def.headers['Sec-WebSocket-Key'] = 'YQ=='  # BASE64 of 'a'
    342         bad_cases.append(
    343             ('Wrong Sec-WebSocket-Key (decoded value is not 16 octets long)',
    344              request_def, 400, True))
    345 
    346         request_def = _create_good_request_def()
    347         del request_def.headers['Sec-WebSocket-Version']
    348         bad_cases.append(('Missing Sec-WebSocket-Version', request_def, None,
    349                           True))
    350 
    351         request_def = _create_good_request_def()
    352         request_def.headers['Sec-WebSocket-Version'] = '3'
    353         bad_cases.append(('Wrong Sec-WebSocket-Version', request_def, None,
    354                           False))
    355 
    356         for (case_name, request_def, expected_status,
    357              expect_handshake_exception) in bad_cases:
    358             request = _create_request(request_def)
    359             handshaker = Handshaker(request, mock.MockDispatcher())
    360             try:
    361                 handshaker.do_handshake()
    362                 self.fail('No exception thrown for \'%s\' case' % case_name)
    363             except HandshakeException, e:
    364                 self.assertTrue(expect_handshake_exception)
    365                 self.assertEqual(expected_status, e.status)
    366             except VersionException, e:
    367                 self.assertFalse(expect_handshake_exception)
    368 
    369 
    370 if __name__ == '__main__':
    371     unittest.main()
    372 
    373 
    374 # vi:sts=4 sw=4 et
    375