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