examples/bench: Add SSL mode

Author: 1st1

examples/bench/echoclient.py

@@ -5,6 +5,7 @@
 import argparse
 import concurrent.futures
 import socket
+import ssl
 import time
 
 
@@ -22,8 +23,17 @@
                         help='number of workers')
     parser.add_argument('--addr', default='127.0.0.1:25000', type=str,
                         help='address:port of echoserver')
+    parser.add_argument('--ssl', default=False, action='store_true')
     args = parser.parse_args()
 
+    client_context = None
+    if args.ssl:
+        print('with SSL')
+        client_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+        if hasattr(client_context, 'check_hostname'):
+            client_context.check_hostname = False
+        client_context.verify_mode = ssl.CERT_NONE
+
     unix = False
     if args.addr.startswith('file:'):
         unix = True
@@ -56,7 +66,11 @@ def run_test(n):
         except (OSError, NameError):
             pass
 
+        if client_context:
+            sock = client_context.wrap_socket(sock)
+
         sock.connect(addr)
+
         while n > 0:
             sock.sendall(msg)
             nrecv = 0

examples/bench/echoserver.py

@@ -2,7 +2,9 @@
 import asyncio
 import gc
 import os.path
+import pathlib
 import socket
+import ssl
 
 
 PRINT = 0
@@ -86,6 +88,7 @@ async def print_debug(loop):
     parser.add_argument('--proto', default=False, action='store_true')
     parser.add_argument('--addr', default='127.0.0.1:25000', type=str)
     parser.add_argument('--print', default=False, action='store_true')
+    parser.add_argument('--ssl', default=False, action='store_true')
     args = parser.parse_args()
 
     if args.uvloop:
@@ -119,6 +122,19 @@ async def print_debug(loop):
 
     print('serving on: {}'.format(addr))
 
+    server_context = None
+    if args.ssl:
+        print('with SSL')
+        server_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
+        server_context.load_cert_chain(
+            (pathlib.Path(__file__).parent.parent.parent /
+                'tests' / 'certs' / 'ssl_cert.pem'),
+            (pathlib.Path(__file__).parent.parent.parent /
+                'tests' / 'certs' / 'ssl_key.pem'))
+        if hasattr(server_context, 'check_hostname'):
+            server_context.check_hostname = False
+        server_context.verify_mode = ssl.CERT_NONE
+
     if args.streams:
         if args.proto:
             print('cannot use --stream and --proto simultaneously')
@@ -127,10 +143,12 @@ async def print_debug(loop):
         print('using asyncio/streams')
         if unix:
             coro = asyncio.start_unix_server(echo_client_streams,
-                                             addr, loop=loop)
+                                             addr, loop=loop,
+                                             ssl=server_context)
         else:
             coro = asyncio.start_server(echo_client_streams,
-                                        *addr, loop=loop)
+                                        *addr, loop=loop,
+                                        ssl=server_context)
         srv = loop.run_until_complete(coro)
     elif args.proto:
         if args.streams:
@@ -139,11 +157,17 @@ async def print_debug(loop):
 
         print('using simple protocol')
         if unix:
-            coro = loop.create_unix_server(EchoProtocol, addr)
+            coro = loop.create_unix_server(EchoProtocol, addr,
+                                           ssl=server_context)
         else:
-            coro = loop.create_server(EchoProtocol, *addr)
+            coro = loop.create_server(EchoProtocol, *addr,
+                                      ssl=server_context)
         srv = loop.run_until_complete(coro)
     else:
+        if args.ssl:
+            print('cannot use SSL for loop.sock_* methods')
+            exit(1)
+
         print('using sock_recv/sock_sendall')
         loop.create_task(echo_server(loop, addr, unix))
     try: