Support SSL of RPC server, client and client_context.

tests/test_rpc.py

@@ -6,18 +6,21 @@
 import multiprocessing
 import socket
 import time
+import ssl
 
 import pytest
 
 import thriftpy
 thriftpy.install_import_hook()  # noqa
 
+from thriftpy._compat import PY3
 from thriftpy.rpc import make_server, client_context
 
 
 addressbook = thriftpy.load(os.path.join(os.path.dirname(__file__),
                                          "addressbook.thrift"))
 unix_sock = "/tmp/thriftpy_test.sock"
+SSL_PORT = 50441
 
 
 class Dispatcher(object):
@@ -85,6 +88,22 @@ def fin():
     request.addfinalizer(fin)
 
 
+@pytest.fixture(scope="module")
+def ssl_server(request):
+    ssl_server = make_server(addressbook.AddressBookService, Dispatcher(),
+                             host='localhost', port=SSL_PORT,
+                             certfile="ssl/server.pem")
+    ps = multiprocessing.Process(target=ssl_server.serve)
+    ps.start()
+
+    time.sleep(0.1)
+
+    def fin():
+        if ps.is_alive():
+            ps.terminate()
+    request.addfinalizer(fin)
+
+
 @pytest.fixture(scope="module")
 def person():
     phone1 = addressbook.PhoneNumber()
@@ -110,45 +129,99 @@ def client(timeout=3000):
                           unix_socket=unix_sock, timeout=timeout)
 
 
+def ssl_client(timeout=3000):
+    return client_context(addressbook.AddressBookService,
+                          host='localhost', port=SSL_PORT,
+                          timeout=timeout,
+                          cafile="ssl/CA.pem", certfile="ssl/client.crt",
+                          keyfile="ssl/client.key")
+
+
 def test_void_api(server):
     with client() as c:
         assert c.ping() is None
 
 
+def test_void_api_with_ssl(ssl_server):
+    with ssl_client() as c:
+        assert c.ping() is None
+
+
 def test_string_api(server):
     with client() as c:
         assert c.hello("world") == "hello world"
 
 
+def test_string_api_with_ssl(ssl_server):
+    with ssl_client() as c:
+        assert c.hello("world") == "hello world"
+
+
 def test_huge_res(server):
     with client() as c:
         big_str = "world" * 100000
         assert c.hello(big_str) == "hello " + big_str
 
 
+def test_huge_res_with_ssl(ssl_server):
+    with ssl_client() as c:
+        big_str = "world" * 100000
+        assert c.hello(big_str) == "hello " + big_str
+
+
 def test_tstruct_req(person):
     with client() as c:
         assert c.add(person) is True
 
 
+def test_tstruct_req_with_ssl(person):
+    with ssl_client() as c:
+        assert c.add(person) is True
+
+
 def test_tstruct_res(person):
     with client() as c:
         assert person == c.get("Alice")
 
 
+def test_tstruct_res_with_ssl(person):
+    with ssl_client() as c:
+        assert person == c.get("Alice")
+
+
 def test_complex_tstruct():
     with client() as c:
         assert len(c.get_phonenumbers("Alice", 0)) == 0
         assert len(c.get_phonenumbers("Alice", 1000)) == 1000
 
 
+def test_complex_tstruct_with_ssl():
+    with ssl_client() as c:
+        assert len(c.get_phonenumbers("Alice", 0)) == 0
+        assert len(c.get_phonenumbers("Alice", 1000)) == 1000
+
+
 def test_exception():
     with pytest.raises(addressbook.PersonNotExistsError):
         with client() as c:
             c.remove("Bob")
 
 
+def test_exception_iwth_ssl():
+    with pytest.raises(addressbook.PersonNotExistsError):
+        with ssl_client() as c:
+            c.remove("Bob")
+
+
 def test_client_timeout():
     with pytest.raises(socket.timeout):
         with client(timeout=500) as c:
             c.sleep(1000)
+
+
+def test_ssl_client_timeout():
+    # SSL socket timeout raises socket.timeout since Python 3.2.
+    # http://bugs.python.org/issue10272
+    with pytest.raises(socket.timeout if PY3 else ssl.SSLError):
+        with ssl_client(timeout=500) as c:
+            c.sleep(1000)

thriftpy/rpc.py

@@ -11,18 +11,29 @@
 from thriftpy.transport import (
     TBufferedTransportFactory,
     TServerSocket,
+    TSSLServerSocket,
     TSocket,
+    TSSLSocket,
 )
 
 
 def make_client(service, host="localhost", port=9090, unix_socket=None,
                 proto_factory=TBinaryProtocolFactory(),
                 trans_factory=TBufferedTransportFactory(),
-                timeout=None):
+                timeout=None,
+                cafile=None, ssl_context=None, certfile=None, keyfile=None):
     if unix_socket:
         socket = TSocket(unix_socket=unix_socket)
+        if certfile:
+            warnings.warn("SSL only works with host:port, not unix_socket.")
     elif host and port:
-        socket = TSocket(host, port, socket_timeout=timeout)
+        if cafile or ssl_context:
+            socket = TSSLSocket(host, port, socket_timeout=timeout,
+                                cafile=cafile,
+                                certfile=certfile, keyfile=keyfile,
+                                ssl_context=ssl_context)
+        else:
+            socket = TSocket(host, port, socket_timeout=timeout)
     else:
         raise ValueError("Either host/port or unix_socket must be provided.")
 
@@ -35,12 +46,20 @@ def make_client(service, host="localhost", port=9090, unix_socket=None,
 def make_server(service, handler,
                 host="localhost", port=9090, unix_socket=None,
                 proto_factory=TBinaryProtocolFactory(),
-                trans_factory=TBufferedTransportFactory()):
+                trans_factory=TBufferedTransportFactory(),
+                certfile=None):
     processor = TProcessor(service, handler)
+
     if unix_socket:
         server_socket = TServerSocket(unix_socket=unix_socket)
+        if certfile:
+            warnings.warn("SSL only works with host:port, not unix_socket.")
     elif host and port:
-        server_socket = TServerSocket(host=host, port=port)
+        if certfile:
+            server_socket = TSSLServerSocket(host=host, port=port,
+                                             certfile=certfile)
+        else:
+            server_socket = TServerSocket(host=host, port=port)
     else:
         raise ValueError("Either host/port or unix_socket must be provided.")
 
@@ -54,7 +73,8 @@ def make_server(service, handler,
 def client_context(service, host="localhost", port=9090, unix_socket=None,
                    proto_factory=TBinaryProtocolFactory(),
                    trans_factory=TBufferedTransportFactory(),
-                   timeout=3000, socket_timeout=3000, connect_timeout=None):
+                   timeout=3000, socket_timeout=3000, connect_timeout=None,
+                   cafile=None, ssl_context=None, certfile=None, keyfile=None):
     if timeout:
         warnings.warn("`timeout` deprecated, use `socket_timeout` and "
                       "`connect_timeout` instead.")
@@ -64,10 +84,20 @@ def client_context(service, host="localhost", port=9090, unix_socket=None,
         socket = TSocket(unix_socket=unix_socket,
                          connect_timeout=connect_timeout,
                          socket_timeout=socket_timeout)
+        if certfile:
+            warnings.warn("SSL only works with host:port, not unix_socket.")
     elif host and port:
-        socket = TSocket(host, port,
-                         connect_timeout=connect_timeout,
-                         socket_timeout=socket_timeout)
+        if cafile or ssl_context:
+            socket = TSSLSocket(host, port,
+                                connect_timeout=connect_timeout,
+                                socket_timeout=socket_timeout,
+                                cafile=cafile,
+                                certfile=certfile, keyfile=keyfile,
+                                ssl_context=ssl_context)
+        else:
+            socket = TSocket(host, port,
+                             connect_timeout=connect_timeout,
+                             socket_timeout=socket_timeout)
     else:
         raise ValueError("Either host/port or unix_socket must be provided.")