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
36import unittest
37
38import set_sys_path  # Update sys.path to locate mod_pywebsocket module.
39
40from mod_pywebsocket.handshake import draft75 as handshake
41from test import mock
42
43
44_GOOD_REQUEST = (
45    80,
46    '/demo',
47    {
48        'Upgrade': 'WebSocket',
49        'Connection': 'Upgrade',
50        'Host': 'example.com',
51        'Origin': 'http://example.com',
52        'WebSocket-Protocol': 'sample',
53    })
54
55_GOOD_RESPONSE_DEFAULT_PORT = (
56    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n'
57    'Upgrade: WebSocket\r\n'
58    'Connection: Upgrade\r\n'
59    'WebSocket-Origin: http://example.com\r\n'
60    'WebSocket-Location: ws://example.com/demo\r\n'
61    'WebSocket-Protocol: sample\r\n'
62    '\r\n')
63
64_GOOD_RESPONSE_SECURE = (
65    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n'
66    'Upgrade: WebSocket\r\n'
67    'Connection: Upgrade\r\n'
68    'WebSocket-Origin: http://example.com\r\n'
69    'WebSocket-Location: wss://example.com/demo\r\n'
70    'WebSocket-Protocol: sample\r\n'
71    '\r\n')
72
73_GOOD_REQUEST_NONDEFAULT_PORT = (
74    8081,
75    '/demo',
76    {
77        'Upgrade': 'WebSocket',
78        'Connection': 'Upgrade',
79        'Host': 'example.com:8081',
80        'Origin': 'http://example.com',
81        'WebSocket-Protocol': 'sample',
82    })
83
84_GOOD_RESPONSE_NONDEFAULT_PORT = (
85    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n'
86    'Upgrade: WebSocket\r\n'
87    'Connection: Upgrade\r\n'
88    'WebSocket-Origin: http://example.com\r\n'
89    'WebSocket-Location: ws://example.com:8081/demo\r\n'
90    'WebSocket-Protocol: sample\r\n'
91    '\r\n')
92
93_GOOD_RESPONSE_SECURE_NONDEF = (
94    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n'
95    'Upgrade: WebSocket\r\n'
96    'Connection: Upgrade\r\n'
97    'WebSocket-Origin: http://example.com\r\n'
98    'WebSocket-Location: wss://example.com:8081/demo\r\n'
99    'WebSocket-Protocol: sample\r\n'
100    '\r\n')
101
102_GOOD_REQUEST_NO_PROTOCOL = (
103    80,
104    '/demo',
105    {
106        'Upgrade': 'WebSocket',
107        'Connection': 'Upgrade',
108        'Host': 'example.com',
109        'Origin': 'http://example.com',
110    })
111
112_GOOD_RESPONSE_NO_PROTOCOL = (
113    'HTTP/1.1 101 Web Socket Protocol Handshake\r\n'
114    'Upgrade: WebSocket\r\n'
115    'Connection: Upgrade\r\n'
116    'WebSocket-Origin: http://example.com\r\n'
117    'WebSocket-Location: ws://example.com/demo\r\n'
118    '\r\n')
119
120_GOOD_REQUEST_WITH_OPTIONAL_HEADERS = (
121    80,
122    '/demo',
123    {
124        'Upgrade': 'WebSocket',
125        'Connection': 'Upgrade',
126        'Host': 'example.com',
127        'Origin': 'http://example.com',
128        'WebSocket-Protocol': 'sample',
129        'AKey': 'AValue',
130        'EmptyValue': '',
131    })
132
133_BAD_REQUESTS = (
134    (  # HTTP request
135        80,
136        '/demo',
137        {
138            'Host': 'www.google.com',
139            'User-Agent': 'Mozilla/5.0 (Macintosh; U; Intel Mac OS X 10.5;'
140                          ' en-US; rv:1.9.1.3) Gecko/20090824 Firefox/3.5.3'
141                          ' GTB6 GTBA',
142            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,'
143                      '*/*;q=0.8',
144            'Accept-Language': 'en-us,en;q=0.5',
145            'Accept-Encoding': 'gzip,deflate',
146            'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
147            'Keep-Alive': '300',
148            'Connection': 'keep-alive',
149        }),
150    (  # Missing Upgrade
151        80,
152        '/demo',
153        {
154            'Connection': 'Upgrade',
155            'Host': 'example.com',
156            'Origin': 'http://example.com',
157            'WebSocket-Protocol': 'sample',
158        }),
159    (  # Wrong Upgrade
160        80,
161        '/demo',
162        {
163            'Upgrade': 'NonWebSocket',
164            'Connection': 'Upgrade',
165            'Host': 'example.com',
166            'Origin': 'http://example.com',
167            'WebSocket-Protocol': 'sample',
168        }),
169    (  # Empty WebSocket-Protocol
170        80,
171        '/demo',
172        {
173            'Upgrade': 'WebSocket',
174            'Connection': 'Upgrade',
175            'Host': 'example.com',
176            'Origin': 'http://example.com',
177            'WebSocket-Protocol': '',
178        }),
179    (  # Wrong port number format
180        80,
181        '/demo',
182        {
183            'Upgrade': 'WebSocket',
184            'Connection': 'Upgrade',
185            'Host': 'example.com:0x50',
186            'Origin': 'http://example.com',
187            'WebSocket-Protocol': 'sample',
188        }),
189    (  # Header/connection port mismatch
190        8080,
191        '/demo',
192        {
193            'Upgrade': 'WebSocket',
194            'Connection': 'Upgrade',
195            'Host': 'example.com',
196            'Origin': 'http://example.com',
197            'WebSocket-Protocol': 'sample',
198        }),
199    (  # Illegal WebSocket-Protocol
200        80,
201        '/demo',
202        {
203            'Upgrade': 'WebSocket',
204            'Connection': 'Upgrade',
205            'Host': 'example.com',
206            'Origin': 'http://example.com',
207            'WebSocket-Protocol': 'illegal\x09protocol',
208        }))
209
210_STRICTLY_GOOD_REQUESTS = (
211    (
212        'GET /demo HTTP/1.1\r\n',
213        'Upgrade: WebSocket\r\n',
214        'Connection: Upgrade\r\n',
215        'Host: example.com\r\n',
216        'Origin: http://example.com\r\n',
217        '\r\n',
218    ),
219    (  # WebSocket-Protocol
220        'GET /demo HTTP/1.1\r\n',
221        'Upgrade: WebSocket\r\n',
222        'Connection: Upgrade\r\n',
223        'Host: example.com\r\n',
224        'Origin: http://example.com\r\n',
225        'WebSocket-Protocol: sample\r\n',
226        '\r\n',
227    ),
228    (  # WebSocket-Protocol and Cookie
229        'GET /demo HTTP/1.1\r\n',
230        'Upgrade: WebSocket\r\n',
231        'Connection: Upgrade\r\n',
232        'Host: example.com\r\n',
233        'Origin: http://example.com\r\n',
234        'WebSocket-Protocol: sample\r\n',
235        'Cookie: xyz\r\n'
236        '\r\n',
237    ),
238    (  # Cookie
239        'GET /demo HTTP/1.1\r\n',
240        'Upgrade: WebSocket\r\n',
241        'Connection: Upgrade\r\n',
242        'Host: example.com\r\n',
243        'Origin: http://example.com\r\n',
244        'Cookie: abc/xyz\r\n'
245        'Cookie2: $Version=1\r\n'
246        'Cookie: abc\r\n'
247        '\r\n',
248    ),
249    (
250        'GET / HTTP/1.1\r\n',
251        'Upgrade: WebSocket\r\n',
252        'Connection: Upgrade\r\n',
253        'Host: example.com\r\n',
254        'Origin: http://example.com\r\n',
255        '\r\n',
256    ),
257)
258
259_NOT_STRICTLY_GOOD_REQUESTS = (
260    (  # Extra space after GET
261        'GET  /demo HTTP/1.1\r\n',
262        'Upgrade: WebSocket\r\n',
263        'Connection: Upgrade\r\n',
264        'Host: example.com\r\n',
265        'Origin: http://example.com\r\n',
266        '\r\n',
267    ),
268    (  # Resource name doesn't stat with '/'
269        'GET demo HTTP/1.1\r\n',
270        'Upgrade: WebSocket\r\n',
271        'Connection: Upgrade\r\n',
272        'Host: example.com\r\n',
273        'Origin: http://example.com\r\n',
274        '\r\n',
275    ),
276    (  # No space after :
277        'GET /demo HTTP/1.1\r\n',
278        'Upgrade:WebSocket\r\n',
279        'Connection: Upgrade\r\n',
280        'Host: example.com\r\n',
281        'Origin: http://example.com\r\n',
282        '\r\n',
283    ),
284    (  # Lower case Upgrade header
285        'GET /demo HTTP/1.1\r\n',
286        'upgrade: WebSocket\r\n',
287        'Connection: Upgrade\r\n',
288        'Host: example.com\r\n',
289        'Origin: http://example.com\r\n',
290        '\r\n',
291    ),
292    (  # Connection comes before Upgrade
293        'GET /demo HTTP/1.1\r\n',
294        'Connection: Upgrade\r\n',
295        'Upgrade: WebSocket\r\n',
296        'Host: example.com\r\n',
297        'Origin: http://example.com\r\n',
298        '\r\n',
299    ),
300    (  # Origin comes before Host
301        'GET /demo HTTP/1.1\r\n',
302        'Upgrade: WebSocket\r\n',
303        'Connection: Upgrade\r\n',
304        'Origin: http://example.com\r\n',
305        'Host: example.com\r\n',
306        '\r\n',
307    ),
308    (  # Host continued to the next line
309        'GET /demo HTTP/1.1\r\n',
310        'Upgrade: WebSocket\r\n',
311        'Connection: Upgrade\r\n',
312        'Host: example\r\n',
313        ' .com\r\n',
314        'Origin: http://example.com\r\n',
315        '\r\n',
316    ),
317    ( # Cookie comes before WebSocket-Protocol
318        'GET /demo HTTP/1.1\r\n',
319        'Upgrade: WebSocket\r\n',
320        'Connection: Upgrade\r\n',
321        'Host: example.com\r\n',
322        'Origin: http://example.com\r\n',
323        'Cookie: xyz\r\n'
324        'WebSocket-Protocol: sample\r\n',
325        '\r\n',
326    ),
327    (  # Unknown header
328        'GET /demo HTTP/1.1\r\n',
329        'Upgrade: WebSocket\r\n',
330        'Connection: Upgrade\r\n',
331        'Host: example.com\r\n',
332        'Origin: http://example.com\r\n',
333        'Content-Type: text/html\r\n'
334        '\r\n',
335    ),
336    (  # Cookie with continuation lines
337        'GET /demo HTTP/1.1\r\n',
338        'Upgrade: WebSocket\r\n',
339        'Connection: Upgrade\r\n',
340        'Host: example.com\r\n',
341        'Origin: http://example.com\r\n',
342        'Cookie: xyz\r\n',
343        ' abc\r\n',
344        ' defg\r\n',
345        '\r\n',
346    ),
347    (  # Wrong-case cookie
348        'GET /demo HTTP/1.1\r\n',
349        'Upgrade: WebSocket\r\n',
350        'Connection: Upgrade\r\n',
351        'Host: example.com\r\n',
352        'Origin: http://example.com\r\n',
353        'cookie: abc/xyz\r\n'
354        '\r\n',
355    ),
356    (  # Cookie, no space after colon
357        'GET /demo HTTP/1.1\r\n',
358        'Upgrade: WebSocket\r\n',
359        'Connection: Upgrade\r\n',
360        'Host: example.com\r\n',
361        'Origin: http://example.com\r\n',
362        'Cookie:abc/xyz\r\n'
363        '\r\n',
364    ),
365)
366
367
368def _create_request(request_def):
369    conn = mock.MockConn('')
370    conn.local_addr = ('0.0.0.0', request_def[0])
371    return mock.MockRequest(
372            uri=request_def[1],
373            headers_in=request_def[2],
374            connection=conn)
375
376
377def _create_get_memorized_lines(lines):
378    """Creates a function that returns the given string."""
379
380    def get_memorized_lines():
381        return lines
382    return get_memorized_lines
383
384
385def _create_requests_with_lines(request_lines_set):
386    requests = []
387    for lines in request_lines_set:
388        request = _create_request(_GOOD_REQUEST)
389        request.connection.get_memorized_lines = _create_get_memorized_lines(
390                lines)
391        requests.append(request)
392    return requests
393
394
395class HandshakerTest(unittest.TestCase):
396    """A unittest for draft75 module."""
397
398    def test_good_request_default_port(self):
399        request = _create_request(_GOOD_REQUEST)
400        handshaker = handshake.Handshaker(request,
401                                          mock.MockDispatcher())
402        handshaker.do_handshake()
403        self.assertEqual(_GOOD_RESPONSE_DEFAULT_PORT,
404                         request.connection.written_data())
405        self.assertEqual('/demo', request.ws_resource)
406        self.assertEqual('http://example.com', request.ws_origin)
407        self.assertEqual('ws://example.com/demo', request.ws_location)
408        self.assertEqual('sample', request.ws_protocol)
409
410    def test_good_request_secure_default_port(self):
411        request = _create_request(_GOOD_REQUEST)
412        request.connection.local_addr = ('0.0.0.0', 443)
413        request.is_https_ = True
414        handshaker = handshake.Handshaker(request,
415                                          mock.MockDispatcher())
416        handshaker.do_handshake()
417        self.assertEqual(_GOOD_RESPONSE_SECURE,
418                         request.connection.written_data())
419        self.assertEqual('sample', request.ws_protocol)
420
421    def test_good_request_nondefault_port(self):
422        request = _create_request(_GOOD_REQUEST_NONDEFAULT_PORT)
423        handshaker = handshake.Handshaker(request,
424                                          mock.MockDispatcher())
425        handshaker.do_handshake()
426        self.assertEqual(_GOOD_RESPONSE_NONDEFAULT_PORT,
427                         request.connection.written_data())
428        self.assertEqual('sample', request.ws_protocol)
429
430    def test_good_request_secure_non_default_port(self):
431        request = _create_request(_GOOD_REQUEST_NONDEFAULT_PORT)
432        request.is_https_ = True
433        handshaker = handshake.Handshaker(request,
434                                          mock.MockDispatcher())
435        handshaker.do_handshake()
436        self.assertEqual(_GOOD_RESPONSE_SECURE_NONDEF,
437                         request.connection.written_data())
438        self.assertEqual('sample', request.ws_protocol)
439
440    def test_good_request_default_no_protocol(self):
441        request = _create_request(_GOOD_REQUEST_NO_PROTOCOL)
442        handshaker = handshake.Handshaker(request,
443                                          mock.MockDispatcher())
444        handshaker.do_handshake()
445        self.assertEqual(_GOOD_RESPONSE_NO_PROTOCOL,
446                         request.connection.written_data())
447        self.assertEqual(None, request.ws_protocol)
448
449    def test_good_request_optional_headers(self):
450        request = _create_request(_GOOD_REQUEST_WITH_OPTIONAL_HEADERS)
451        handshaker = handshake.Handshaker(request,
452                                          mock.MockDispatcher())
453        handshaker.do_handshake()
454        self.assertEqual('AValue',
455                         request.headers_in['AKey'])
456        self.assertEqual('',
457                         request.headers_in['EmptyValue'])
458
459    def test_bad_requests(self):
460        for request in map(_create_request, _BAD_REQUESTS):
461            handshaker = handshake.Handshaker(request,
462                                              mock.MockDispatcher())
463            self.assertRaises(
464                handshake.HandshakeException, handshaker.do_handshake)
465
466    def test_strictly_good_requests(self):
467        for request in _create_requests_with_lines(_STRICTLY_GOOD_REQUESTS):
468            strict_handshaker = handshake.Handshaker(request,
469                                                     mock.MockDispatcher(),
470                                                     True)
471            strict_handshaker.do_handshake()
472
473    def test_not_strictly_good_requests(self):
474        for request in _create_requests_with_lines(
475            _NOT_STRICTLY_GOOD_REQUESTS):
476            strict_handshaker = handshake.Handshaker(request,
477                                                     mock.MockDispatcher(),
478                                                     True)
479            self.assertRaises(handshake.HandshakeException,
480                              strict_handshaker.do_handshake)
481
482
483if __name__ == '__main__':
484    unittest.main()
485
486
487# vi:sts=4 sw=4 et
488