From dc70cc38123326a916203c513daa0a1b35b15fdf Mon Sep 17 00:00:00 2001 From: Ryan Williams Date: Thu, 16 Sep 2010 11:42:03 -0700 Subject: [PATCH] Support for SSL websockets, which also happens to improve our SSL support in WSGI generally. Fixes #62. --- eventlet/green/ssl.py | 4 +++- eventlet/websocket.py | 6 ++++- eventlet/wsgi.py | 6 ++++- examples/websocket_chat.py | 2 ++ tests/__init__.py | 3 +++ tests/ssl_test.py | 5 +---- tests/websocket_test.py | 46 +++++++++++++++++++++++++++++++++++++- 7 files changed, 64 insertions(+), 8 deletions(-) diff --git a/eventlet/green/ssl.py b/eventlet/green/ssl.py index 177cf68..c9ba7f8 100644 --- a/eventlet/green/ssl.py +++ b/eventlet/green/ssl.py @@ -300,7 +300,9 @@ class GreenSSLSocket(__ssl.SSLSocket): do_handshake_on_connect=self.do_handshake_on_connect, suppress_ragged_eofs=self.suppress_ragged_eofs) return (new_ssl, addr) - + + def dup(self): + raise NotImplementedError("Can't dup an ssl object") SSLSocket = GreenSSLSocket diff --git a/eventlet/websocket.py b/eventlet/websocket.py index 51a0665..3c6d071 100644 --- a/eventlet/websocket.py +++ b/eventlet/websocket.py @@ -71,7 +71,11 @@ class WebSocketWSGI(object): response = md5(key).digest() # Start building the response - location = 'ws://%s%s%s' % ( + scheme = 'ws' + if environ.get('wsgi.url_scheme') == 'https': + scheme = 'wss' + location = '%s://%s%s%s' % ( + scheme, environ.get('HTTP_HOST'), environ.get('SCRIPT_NAME'), environ.get('PATH_INFO') diff --git a/eventlet/wsgi.py b/eventlet/wsgi.py index 5c74386..c49f132 100644 --- a/eventlet/wsgi.py +++ b/eventlet/wsgi.py @@ -160,7 +160,7 @@ class Input(object): return iter(self.read()) def get_socket(self): - return self.rfile._sock.dup() + return self.rfile._sock class HttpProtocol(BaseHTTPServer.BaseHTTPRequestHandler): @@ -498,6 +498,10 @@ class Server(BaseHTTPServer.HTTPServer): 'wsgi.run_once': False, 'wsgi.url_scheme': 'http', } + # detect secure socket + if hasattr(self.socket, 'do_handshake'): + d['wsgi.url_scheme'] = 'https' + d['HTTPS'] = 'on' if self.environ is not None: d.update(self.environ) return d diff --git a/examples/websocket_chat.py b/examples/websocket_chat.py index 7f7e3ea..bf182e1 100644 --- a/examples/websocket_chat.py +++ b/examples/websocket_chat.py @@ -1,3 +1,5 @@ +import os + import eventlet from eventlet import wsgi from eventlet import websocket diff --git a/tests/__init__.py b/tests/__init__.py index cf9c58a..f5225d3 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -224,3 +224,6 @@ def get_database_auth(): except IOError: pass return retval + +certificate_file = os.path.join(os.path.dirname(__file__), 'test_server.crt') +private_key_file = os.path.join(os.path.dirname(__file__), 'test_server.key') diff --git a/tests/ssl_test.py b/tests/ssl_test.py index 7c5ad3d..1081366 100644 --- a/tests/ssl_test.py +++ b/tests/ssl_test.py @@ -1,13 +1,10 @@ -from tests import skipped, LimitedTestCase, skip_unless +from tests import skipped, LimitedTestCase, skip_unless, certificate_file, private_key_file from unittest import main import eventlet from eventlet import util, coros, greenio import socket import os -certificate_file = os.path.join(os.path.dirname(__file__), 'test_server.crt') -private_key_file = os.path.join(os.path.dirname(__file__), 'test_server.key') - def listen_ssl_socket(address=('127.0.0.1', 0)): sock = util.wrap_ssl(socket.socket(), certificate_file, private_key_file, True) diff --git a/tests/websocket_test.py b/tests/websocket_test.py index 2e0fa25..db68fb0 100644 --- a/tests/websocket_test.py +++ b/tests/websocket_test.py @@ -8,7 +8,7 @@ from eventlet.websocket import WebSocket, WebSocketWSGI from eventlet import wsgi from eventlet import event -from tests import mock, LimitedTestCase +from tests import mock, LimitedTestCase, certificate_file, private_key_file from tests.wsgi_test import _TestBase @@ -460,6 +460,50 @@ class TestWebSocket(_TestBase): self.assert_(error_detected[0]) +class TestWebSocketSSL(_TestBase): + def set_site(self): + self.site = wsapp + + def test_ssl_sending_messages(self): + s = eventlet.wrap_ssl(eventlet.listen(('localhost', 0)), + certfile=certificate_file, + keyfile=private_key_file, + server_side=True) + self.spawn_server(sock=s) + connect = [ + "GET /echo HTTP/1.1", + "Upgrade: WebSocket", + "Connection: Upgrade", + "Host: localhost:%s" % self.port, + "Origin: http://localhost:%s" % self.port, + "Sec-WebSocket-Protocol: ws", + "Sec-WebSocket-Key1: 4 @1 46546xW%0l 1 5", + "Sec-WebSocket-Key2: 12998 5 Y3 1 .P00", + ] + sock = eventlet.wrap_ssl(eventlet.connect( + ('localhost', self.port))) + + sock.sendall('\r\n'.join(connect) + '\r\n\r\n^n:ds[4U') + first_resp = sock.recv(1024) + # make sure it sets the wss: protocol on the location header + loc_line = [x for x in first_resp.split("\r\n") + if x.lower().startswith('sec-websocket-location')][0] + self.assert_("wss://localhost" in loc_line, + "Expecting wss protocol in location: %s" % loc_line) + sock.sendall('\x00hello\xFF') + result = sock.recv(1024) + self.assertEqual(result, '\x00hello\xff') + sock.sendall('\x00start') + eventlet.sleep(0.001) + sock.sendall(' end\xff') + result = sock.recv(1024) + self.assertEqual(result, '\x00start end\xff') + sock.shutdown(socket.SHUT_RDWR) + sock.close() + eventlet.sleep(0.01) + + + class TestWebSocketObject(LimitedTestCase): def setUp(self):