Enable TLS protected connections to memcache

aiomcache/client.py

@@ -1,6 +1,7 @@
 import functools
 import re
 import sys
+from ssl import SSLContext
 from typing import (Any, Awaitable, Callable, Dict, Generic, Optional, Tuple, TypeVar,
                     Union, overload)
 
@@ -52,6 +53,8 @@ async def wrapper(self: _Client, *args: _P.args,  # type: ignore[misc]
 class FlagClient(Generic[_T]):
     def __init__(self, host: str, port: int = 11211, *,
                  pool_size: int = 2, pool_minsize: Optional[int] = None,
+                 ssl: Union[bool, SSLContext] = False,
+                 ssl_handshake_timeout: Union[float, None] = None,
                  get_flag_handler: Optional[_GetFlagHandler[_T]] = None,
                  set_flag_handler: Optional[_SetFlagHandler[_T]] = None):
         """
@@ -61,6 +64,12 @@ def __init__(self, host: str, port: int = 11211, *,
         :param port: memcached port
         :param pool_size: max connection pool size
         :param pool_minsize: min connection pool size
+        :param ssl: whether to use TLS against the memcache server. If ssl
+            is a ssl.SSLContext object, this context is used to create the
+            transport; if ssl is True, a default context returned from
+            ssl.create_default_context() is used
+        :param ssl_handshake_timeout: time in seconds to wait for the TLS
+            handshake to complete, if using TLS
         :param get_flag_handler: async method to call to convert flagged
             values. Method takes tuple: (value, flags) and should return
             processed value or raise ClientException if not supported.
@@ -72,7 +81,8 @@ def __init__(self, host: str, port: int = 11211, *,
             pool_minsize = pool_size
 
         self._pool = MemcachePool(
-            host, port, minsize=pool_minsize, maxsize=pool_size)
+            host, port, minsize=pool_minsize, maxsize=pool_size, ssl=ssl,
+            ssl_handshake_timeout=ssl_handshake_timeout)
 
         self._get_flag_handler = get_flag_handler
         self._set_flag_handler = set_flag_handler
@@ -493,6 +503,9 @@ async def flush_all(self, conn: Connection) -> None:
 
 class Client(FlagClient[bytes]):
     def __init__(self, host: str, port: int = 11211, *,
-                 pool_size: int = 2, pool_minsize: Optional[int] = None):
+                 pool_size: int = 2, pool_minsize: Optional[int] = None,
+                 ssl: Union[bool, SSLContext] = False,
+                 ssl_handshake_timeout: Union[float, None] = None):
         super().__init__(host, port, pool_size=pool_size, pool_minsize=pool_minsize,
+                         ssl=ssl, ssl_handshake_timeout=ssl_handshake_timeout,
                          get_flag_handler=None, set_flag_handler=None)

aiomcache/pool.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import NamedTuple, Optional, Set
+from ssl import SSLContext
+from typing import NamedTuple, Optional, Set, Union
 
 __all__ = ['MemcachePool']
 
@@ -10,11 +11,15 @@ class Connection(NamedTuple):
 
 
 class MemcachePool:
-    def __init__(self, host: str, port: int, *, minsize: int, maxsize: int):
+    def __init__(self, host: str, port: int, *, minsize: int, maxsize: int,
+                 ssl: Union[bool, SSLContext] = False,
+                 ssl_handshake_timeout: Union[float, None] = None):
         self._host = host
         self._port = port
         self._minsize = minsize
         self._maxsize = maxsize
+        self._ssl = ssl
+        self._ssl_handshake_timeout = ssl_handshake_timeout
         self._pool: asyncio.Queue[Connection] = asyncio.Queue()
         self._in_use: Set[Connection] = set()
 
@@ -66,7 +71,8 @@ def release(self, conn: Connection) -> None:
     async def _create_new_conn(self) -> Optional[Connection]:
         if self.size() < self._maxsize:
             reader, writer = await asyncio.open_connection(
-                self._host, self._port)
+                self._host, self._port, ssl=self._ssl,
+                ssl_handshake_timeout=self._ssl_handshake_timeout)
             if self.size() < self._maxsize:
                 return Connection(reader, writer)
             else:

tests/ssl_test.py

@@ -0,0 +1,28 @@
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+import ssl
+from aiomcache import Client
+from .conftest import McacheParams
+
+
+@pytest.mark.asyncio
+async def test_ssl_params_forwarded_from_client() -> None:
+        client = Client("host", port=11211, ssl=True, ssl_handshake_timeout=20)
+
+        with mock.patch("asyncio.open_connection", return_value=(MagicMock, MagicMock)) as oc:
+            await client._pool._create_new_conn()
+
+        oc.assert_called_once_with("host", 11211, ssl=True, ssl_handshake_timeout=20)
+
+@pytest.mark.asyncio
+async def test_ssl_client_fails_against_plaintext_server(
+    mcache_params: McacheParams
+) -> None:
+        client = Client(**mcache_params, ssl=True)
+        # If SSL was correctly enabled, this should
+        # fail, since SSL isn't enabled on the memcache
+        # server.
+        with pytest.raises(ssl.SSLError):
+            await client.get(b"key")