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