tests/websocket_test.py

   1 import errno
   2 import socket
   3
   4 import eventlet
   5 from eventlet import event
   6 from eventlet import greenio
   7 from eventlet import wsgi
   8 from eventlet.green import httplib
   9 from eventlet.green import urllib2
  10 from eventlet.websocket import WebSocket, WebSocketWSGI
  11
  12 from tests import mock, LimitedTestCase, certificate_file, private_key_file
  13 from tests import skip_if_no_ssl
  14 from tests.wsgi_test import _TestBase
  15
  16
  17 # demo app
  18 def handle(ws):
  19     if ws.path == '/echo':
  20         while True:
  21             m = ws.wait()
  22             if m is None:
  23                 break
  24             ws.send(m)
  25     elif ws.path == '/range':
  26         for i in range(10):
  27             ws.send("msg %d" % i)
  28             eventlet.sleep(0.01)
  29     elif ws.path == '/error':
  30         # some random socket error that we shouldn't normally get
  31         raise socket.error(errno.ENOTSOCK)
  32     else:
  33         ws.close()
  34
  35 wsapp = WebSocketWSGI(handle)
  36
  37
  38 class TestWebSocket(_TestBase):
  39     TEST_TIMEOUT = 5
  40
  41     def set_site(self):
  42         self.site = wsapp
  43
  44     def test_incorrect_headers(self):
  45         def raiser():
  46             try:
  47                 urllib2.urlopen("http://localhost:%s/echo" % self.port)
  48             except urllib2.HTTPError as e:
  49                 self.assertEqual(e.code, 400)
  50                 raise
  51         self.assertRaises(urllib2.HTTPError, raiser)
  52
  53     def test_incomplete_headers_75(self):
  54         headers = dict(kv.split(': ') for kv in [
  55             "Upgrade: WebSocket",
  56             # NOTE: intentionally no connection header
  57             "Host: localhost:%s" % self.port,
  58             "Origin: http://localhost:%s" % self.port,
  59             "WebSocket-Protocol: ws",
  60         ])
  61         http = httplib.HTTPConnection('localhost', self.port)
  62         http.request("GET", "/echo", headers=headers)
  63         resp = http.getresponse()
  64
  65         self.assertEqual(resp.status, 400)
  66         self.assertEqual(resp.getheader('connection'), 'close')
  67         self.assertEqual(resp.read(), '')
  68
  69     def test_incomplete_headers_76(self):
  70         # First test: Missing Connection:
  71         headers = dict(kv.split(': ') for kv in [
  72             "Upgrade: WebSocket",
  73             # NOTE: intentionally no connection header
  74             "Host: localhost:%s" % self.port,
  75             "Origin: http://localhost:%s" % self.port,
  76             "Sec-WebSocket-Protocol: ws",
  77         ])
  78         http = httplib.HTTPConnection('localhost', self.port)
  79         http.request("GET", "/echo", headers=headers)
  80         resp = http.getresponse()
  81
  82         self.assertEqual(resp.status, 400)
  83         self.assertEqual(resp.getheader('connection'), 'close')
  84         self.assertEqual(resp.read(), '')
  85
  86         # Now, miss off key2
  87         headers = dict(kv.split(': ') for kv in [
  88             "Upgrade: WebSocket",
  89             "Connection: Upgrade",
  90             "Host: localhost:%s" % self.port,
  91             "Origin: http://localhost:%s" % self.port,
  92             "Sec-WebSocket-Protocol: ws",
  93             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
  94             # NOTE: Intentionally no Key2 header
  95         ])
  96         http = httplib.HTTPConnection('localhost', self.port)
  97         http.request("GET", "/echo", headers=headers)
  98         resp = http.getresponse()
  99
 100         self.assertEqual(resp.status, 400)
 101         self.assertEqual(resp.getheader('connection'), 'close')
 102         self.assertEqual(resp.read(), '')
 103
 104     def test_correct_upgrade_request_75(self):
 105         connect = [
 106             "GET /echo HTTP/1.1",
 107             "Upgrade: WebSocket",
 108             "Connection: Upgrade",
 109             "Host: localhost:%s" % self.port,
 110             "Origin: http://localhost:%s" % self.port,
 111             "WebSocket-Protocol: ws",
 112         ]
 113         sock = eventlet.connect(
 114             ('localhost', self.port))
 115
 116         sock.sendall('\r\n'.join(connect) + '\r\n\r\n')
 117         result = sock.recv(1024)
 118         # The server responds the correct Websocket handshake
 119         self.assertEqual(result, '\r\n'.join([
 120             'HTTP/1.1 101 Web Socket Protocol Handshake',
 121             'Upgrade: WebSocket',
 122             'Connection: Upgrade',
 123             'WebSocket-Origin: http://localhost:%s' % self.port,
 124             'WebSocket-Location: ws://localhost:%s/echo\r\n\r\n' % self.port,
 125         ]))
 126
 127     def test_correct_upgrade_request_76(self):
 128         connect = [
 129             "GET /echo HTTP/1.1",
 130             "Upgrade: WebSocket",
 131             "Connection: Upgrade",
 132             "Host: localhost:%s" % self.port,
 133             "Origin: http://localhost:%s" % self.port,
 134             "Sec-WebSocket-Protocol: ws",
 135             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 136             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 137         ]
 138         sock = eventlet.connect(
 139             ('localhost', self.port))
 140
 141         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 142         result = sock.recv(1024)
 143         # The server responds the correct Websocket handshake
 144         self.assertEqual(result, '\r\n'.join([
 145             'HTTP/1.1 101 WebSocket Protocol Handshake',
 146             'Upgrade: WebSocket',
 147             'Connection: Upgrade',
 148             'Sec-WebSocket-Origin: http://localhost:%s' % self.port,
 149             'Sec-WebSocket-Protocol: ws',
 150             'Sec-WebSocket-Location: ws://localhost:%s/echo\r\n\r\n8jKS\'y:G*Co,Wxa-' % self.port,
 151         ]))
 152
 153     def test_query_string(self):
 154         # verify that the query string comes out the other side unscathed
 155         connect = [
 156             "GET /echo?query_string HTTP/1.1",
 157             "Upgrade: WebSocket",
 158             "Connection: Upgrade",
 159             "Host: localhost:%s" % self.port,
 160             "Origin: http://localhost:%s" % self.port,
 161             "Sec-WebSocket-Protocol: ws",
 162             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 163             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 164         ]
 165         sock = eventlet.connect(
 166             ('localhost', self.port))
 167
 168         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 169         result = sock.recv(1024)
 170         self.assertEqual(result, '\r\n'.join([
 171             'HTTP/1.1 101 WebSocket Protocol Handshake',
 172             'Upgrade: WebSocket',
 173             'Connection: Upgrade',
 174             'Sec-WebSocket-Origin: http://localhost:%s' % self.port,
 175             'Sec-WebSocket-Protocol: ws',
 176             'Sec-WebSocket-Location: ws://localhost:%s/echo?query_string\r\n\r\n8jKS\'y:G*Co,Wxa-' % self.port,
 177         ]))
 178
 179     def test_empty_query_string(self):
 180         # verify that a single trailing ? doesn't get nuked
 181         connect = [
 182             "GET /echo? HTTP/1.1",
 183             "Upgrade: WebSocket",
 184             "Connection: Upgrade",
 185             "Host: localhost:%s" % self.port,
 186             "Origin: http://localhost:%s" % self.port,
 187             "Sec-WebSocket-Protocol: ws",
 188             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 189             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 190         ]
 191         sock = eventlet.connect(
 192             ('localhost', self.port))
 193
 194         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 195         result = sock.recv(1024)
 196         self.assertEqual(result, '\r\n'.join([
 197             'HTTP/1.1 101 WebSocket Protocol Handshake',
 198             'Upgrade: WebSocket',
 199             'Connection: Upgrade',
 200             'Sec-WebSocket-Origin: http://localhost:%s' % self.port,
 201             'Sec-WebSocket-Protocol: ws',
 202             'Sec-WebSocket-Location: ws://localhost:%s/echo?\r\n\r\n8jKS\'y:G*Co,Wxa-' % self.port,
 203         ]))
 204
 205     def test_sending_messages_to_websocket_75(self):
 206         connect = [
 207             "GET /echo HTTP/1.1",
 208             "Upgrade: WebSocket",
 209             "Connection: Upgrade",
 210             "Host: localhost:%s" % self.port,
 211             "Origin: http://localhost:%s" % self.port,
 212             "WebSocket-Protocol: ws",
 213         ]
 214         sock = eventlet.connect(
 215             ('localhost', self.port))
 216
 217         sock.sendall('\r\n'.join(connect) + '\r\n\r\n')
 218         sock.recv(1024)
 219         sock.sendall('\x00hello\xFF')
 220         result = sock.recv(1024)
 221         self.assertEqual(result, '\x00hello\xff')
 222         sock.sendall('\x00start')
 223         eventlet.sleep(0.001)
 224         sock.sendall(' end\xff')
 225         result = sock.recv(1024)
 226         self.assertEqual(result, '\x00start end\xff')
 227         sock.shutdown(socket.SHUT_RDWR)
 228         sock.close()
 229         eventlet.sleep(0.01)
 230
 231     def test_sending_messages_to_websocket_76(self):
 232         connect = [
 233             "GET /echo HTTP/1.1",
 234             "Upgrade: WebSocket",
 235             "Connection: Upgrade",
 236             "Host: localhost:%s" % self.port,
 237             "Origin: http://localhost:%s" % self.port,
 238             "Sec-WebSocket-Protocol: ws",
 239             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 240             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 241         ]
 242         sock = eventlet.connect(
 243             ('localhost', self.port))
 244
 245         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 246         sock.recv(1024)
 247         sock.sendall('\x00hello\xFF')
 248         result = sock.recv(1024)
 249         self.assertEqual(result, '\x00hello\xff')
 250         sock.sendall('\x00start')
 251         eventlet.sleep(0.001)
 252         sock.sendall(' end\xff')
 253         result = sock.recv(1024)
 254         self.assertEqual(result, '\x00start end\xff')
 255         sock.shutdown(socket.SHUT_RDWR)
 256         sock.close()
 257         eventlet.sleep(0.01)
 258
 259     def test_getting_messages_from_websocket_75(self):
 260         connect = [
 261             "GET /range HTTP/1.1",
 262             "Upgrade: WebSocket",
 263             "Connection: Upgrade",
 264             "Host: localhost:%s" % self.port,
 265             "Origin: http://localhost:%s" % self.port,
 266             "WebSocket-Protocol: ws",
 267         ]
 268         sock = eventlet.connect(
 269             ('localhost', self.port))
 270
 271         sock.sendall('\r\n'.join(connect) + '\r\n\r\n')
 272         resp = sock.recv(1024)
 273         headers, result = resp.split('\r\n\r\n')
 274         msgs = [result.strip('\x00\xff')]
 275         cnt = 10
 276         while cnt:
 277             msgs.append(sock.recv(20).strip('\x00\xff'))
 278             cnt -= 1
 279         # Last item in msgs is an empty string
 280         self.assertEqual(msgs[:-1], ['msg %d' % i for i in range(10)])
 281
 282     def test_getting_messages_from_websocket_76(self):
 283         connect = [
 284             "GET /range HTTP/1.1",
 285             "Upgrade: WebSocket",
 286             "Connection: Upgrade",
 287             "Host: localhost:%s" % self.port,
 288             "Origin: http://localhost:%s" % self.port,
 289             "Sec-WebSocket-Protocol: ws",
 290             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 291             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 292         ]
 293         sock = eventlet.connect(
 294             ('localhost', self.port))
 295
 296         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 297         resp = sock.recv(1024)
 298         headers, result = resp.split('\r\n\r\n')
 299         msgs = [result[16:].strip('\x00\xff')]
 300         cnt = 10
 301         while cnt:
 302             msgs.append(sock.recv(20).strip('\x00\xff'))
 303             cnt -= 1
 304         # Last item in msgs is an empty string
 305         self.assertEqual(msgs[:-1], ['msg %d' % i for i in range(10)])
 306
 307     def test_breaking_the_connection_75(self):
 308         error_detected = [False]
 309         done_with_request = event.Event()
 310         site = self.site
 311
 312         def error_detector(environ, start_response):
 313             try:
 314                 try:
 315                     return site(environ, start_response)
 316                 except:
 317                     error_detected[0] = True
 318                     raise
 319             finally:
 320                 done_with_request.send(True)
 321         self.site = error_detector
 322         self.spawn_server()
 323         connect = [
 324             "GET /range HTTP/1.1",
 325             "Upgrade: WebSocket",
 326             "Connection: Upgrade",
 327             "Host: localhost:%s" % self.port,
 328             "Origin: http://localhost:%s" % self.port,
 329             "WebSocket-Protocol: ws",
 330         ]
 331         sock = eventlet.connect(
 332             ('localhost', self.port))
 333         sock.sendall('\r\n'.join(connect) + '\r\n\r\n')
 334         sock.recv(1024)  # get the headers
 335         sock.close()  # close while the app is running
 336         done_with_request.wait()
 337         assert not error_detected[0]
 338
 339     def test_breaking_the_connection_76(self):
 340         error_detected = [False]
 341         done_with_request = event.Event()
 342         site = self.site
 343
 344         def error_detector(environ, start_response):
 345             try:
 346                 try:
 347                     return site(environ, start_response)
 348                 except:
 349                     error_detected[0] = True
 350                     raise
 351             finally:
 352                 done_with_request.send(True)
 353         self.site = error_detector
 354         self.spawn_server()
 355         connect = [
 356             "GET /range HTTP/1.1",
 357             "Upgrade: WebSocket",
 358             "Connection: Upgrade",
 359             "Host: localhost:%s" % self.port,
 360             "Origin: http://localhost:%s" % self.port,
 361             "Sec-WebSocket-Protocol: ws",
 362             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 363             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 364         ]
 365         sock = eventlet.connect(
 366             ('localhost', self.port))
 367         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 368         sock.recv(1024)  # get the headers
 369         sock.close()  # close while the app is running
 370         done_with_request.wait()
 371         assert not error_detected[0]
 372
 373     def test_client_closing_connection_76(self):
 374         error_detected = [False]
 375         done_with_request = event.Event()
 376         site = self.site
 377
 378         def error_detector(environ, start_response):
 379             try:
 380                 try:
 381                     return site(environ, start_response)
 382                 except:
 383                     error_detected[0] = True
 384                     raise
 385             finally:
 386                 done_with_request.send(True)
 387         self.site = error_detector
 388         self.spawn_server()
 389         connect = [
 390             "GET /echo HTTP/1.1",
 391             "Upgrade: WebSocket",
 392             "Connection: Upgrade",
 393             "Host: localhost:%s" % self.port,
 394             "Origin: http://localhost:%s" % self.port,
 395             "Sec-WebSocket-Protocol: ws",
 396             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 397             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 398         ]
 399         sock = eventlet.connect(
 400             ('localhost', self.port))
 401         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 402         sock.recv(1024)  # get the headers
 403         sock.sendall('\xff\x00')  # "Close the connection" packet.
 404         done_with_request.wait()
 405         assert not error_detected[0]
 406
 407     def test_client_invalid_packet_76(self):
 408         error_detected = [False]
 409         done_with_request = event.Event()
 410         site = self.site
 411
 412         def error_detector(environ, start_response):
 413             try:
 414                 try:
 415                     return site(environ, start_response)
 416                 except:
 417                     error_detected[0] = True
 418                     raise
 419             finally:
 420                 done_with_request.send(True)
 421         self.site = error_detector
 422         self.spawn_server()
 423         connect = [
 424             "GET /echo HTTP/1.1",
 425             "Upgrade: WebSocket",
 426             "Connection: Upgrade",
 427             "Host: localhost:%s" % self.port,
 428             "Origin: http://localhost:%s" % self.port,
 429             "Sec-WebSocket-Protocol: ws",
 430             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 431             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 432         ]
 433         sock = eventlet.connect(
 434             ('localhost', self.port))
 435         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 436         sock.recv(1024)  # get the headers
 437         sock.sendall('\xef\x00')  # Weird packet.
 438         done_with_request.wait()
 439         assert error_detected[0]
 440
 441     def test_server_closing_connect_76(self):
 442         connect = [
 443             "GET / HTTP/1.1",
 444             "Upgrade: WebSocket",
 445             "Connection: Upgrade",
 446             "Host: localhost:%s" % self.port,
 447             "Origin: http://localhost:%s" % self.port,
 448             "Sec-WebSocket-Protocol: ws",
 449             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 450             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 451         ]
 452         sock = eventlet.connect(
 453             ('localhost', self.port))
 454
 455         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 456         resp = sock.recv(1024)
 457         headers, result = resp.split('\r\n\r\n')
 458         # The remote server should have immediately closed the connection.
 459         self.assertEqual(result[16:], '\xff\x00')
 460
 461     def test_app_socket_errors_75(self):
 462         error_detected = [False]
 463         done_with_request = event.Event()
 464         site = self.site
 465
 466         def error_detector(environ, start_response):
 467             try:
 468                 try:
 469                     return site(environ, start_response)
 470                 except:
 471                     error_detected[0] = True
 472                     raise
 473             finally:
 474                 done_with_request.send(True)
 475         self.site = error_detector
 476         self.spawn_server()
 477         connect = [
 478             "GET /error HTTP/1.1",
 479             "Upgrade: WebSocket",
 480             "Connection: Upgrade",
 481             "Host: localhost:%s" % self.port,
 482             "Origin: http://localhost:%s" % self.port,
 483             "WebSocket-Protocol: ws",
 484         ]
 485         sock = eventlet.connect(
 486             ('localhost', self.port))
 487         sock.sendall('\r\n'.join(connect) + '\r\n\r\n')
 488         sock.recv(1024)
 489         done_with_request.wait()
 490         assert error_detected[0]
 491
 492     def test_app_socket_errors_76(self):
 493         error_detected = [False]
 494         done_with_request = event.Event()
 495         site = self.site
 496
 497         def error_detector(environ, start_response):
 498             try:
 499                 try:
 500                     return site(environ, start_response)
 501                 except:
 502                     error_detected[0] = True
 503                     raise
 504             finally:
 505                 done_with_request.send(True)
 506         self.site = error_detector
 507         self.spawn_server()
 508         connect = [
 509             "GET /error HTTP/1.1",
 510             "Upgrade: WebSocket",
 511             "Connection: Upgrade",
 512             "Host: localhost:%s" % self.port,
 513             "Origin: http://localhost:%s" % self.port,
 514             "Sec-WebSocket-Protocol: ws",
 515             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 516             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 517         ]
 518         sock = eventlet.connect(
 519             ('localhost', self.port))
 520         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 521         sock.recv(1024)
 522         done_with_request.wait()
 523         assert error_detected[0]
 524
 525
 526 class TestWebSocketSSL(_TestBase):
 527     def set_site(self):
 528         self.site = wsapp
 529
 530     @skip_if_no_ssl
 531     def test_ssl_sending_messages(self):
 532         s = eventlet.wrap_ssl(eventlet.listen(('localhost', 0)),
 533                               certfile=certificate_file,
 534                               keyfile=private_key_file,
 535                               server_side=True)
 536         self.spawn_server(sock=s)
 537         connect = [
 538             "GET /echo HTTP/1.1",
 539             "Upgrade: WebSocket",
 540             "Connection: Upgrade",
 541             "Host: localhost:%s" % self.port,
 542             "Origin: http://localhost:%s" % self.port,
 543             "Sec-WebSocket-Protocol: ws",
 544             "Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5",
 545             "Sec-WebSocket-Key2: 12998 5 Y3 1  .P00",
 546         ]
 547         sock = eventlet.wrap_ssl(eventlet.connect(
 548             ('localhost', self.port)))
 549
 550         sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U')
 551         first_resp = sock.recv(1024)
 552         # make sure it sets the wss: protocol on the location header
 553         loc_line = [x for x in first_resp.split("\r\n")
 554                     if x.lower().startswith('sec-websocket-location')][0]
 555         self.assert_("wss://localhost" in loc_line,
 556                      "Expecting wss protocol in location: %s" % loc_line)
 557         sock.sendall('\x00hello\xFF')
 558         result = sock.recv(1024)
 559         self.assertEqual(result, '\x00hello\xff')
 560         sock.sendall('\x00start')
 561         eventlet.sleep(0.001)
 562         sock.sendall(' end\xff')
 563         result = sock.recv(1024)
 564         self.assertEqual(result, '\x00start end\xff')
 565         greenio.shutdown_safe(sock)
 566         sock.close()
 567         eventlet.sleep(0.01)
 568
 569
 570 class TestWebSocketObject(LimitedTestCase):
 571
 572     def setUp(self):
 573         self.mock_socket = s = mock.Mock()
 574         self.environ = env = dict(HTTP_ORIGIN='http://localhost', HTTP_WEBSOCKET_PROTOCOL='ws',
 575                                   PATH_INFO='test')
 576
 577         self.test_ws = WebSocket(s, env)
 578         super(TestWebSocketObject, self).setUp()
 579
 580     def test_recieve(self):
 581         ws = self.test_ws
 582         ws.socket.recv.return_value = '\x00hello\xFF'
 583         self.assertEqual(ws.wait(), 'hello')
 584         self.assertEqual(ws._buf, '')
 585         self.assertEqual(len(ws._msgs), 0)
 586         ws.socket.recv.return_value = ''
 587         self.assertEqual(ws.wait(), None)
 588         self.assertEqual(ws._buf, '')
 589         self.assertEqual(len(ws._msgs), 0)
 590
 591     def test_send_to_ws(self):
 592         ws = self.test_ws
 593         ws.send(u'hello')
 594         assert ws.socket.sendall.called_with("\x00hello\xFF")
 595         ws.send(10)
 596         assert ws.socket.sendall.called_with("\x0010\xFF")
 597
 598     def test_close_ws(self):
 599         ws = self.test_ws
 600         ws.close()
 601         assert ws.socket.shutdown.called_with(True)