Add TLS support for TCP sockets (#276)

2 files changed, 19 insertions, 5 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 972cf4d..0387da4 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -219,7 +219,8 @@ class Client(object):
                  key_prefix=b'',
                  default_noreply=True,
                  allow_unicode_keys=False,
-                 encoding='ascii'):
+                 encoding='ascii',
+                 tls_context=None):
         """
         Constructor.
 
@@ -269,6 +270,7 @@ class Client(object):
         self.default_noreply = default_noreply
         self.allow_unicode_keys = allow_unicode_keys
         self.encoding = encoding
+        self.tls_context = tls_context
 
     def check_key(self, key):
         """Checks key and add key_prefix."""
@@ -281,6 +283,11 @@ class Client(object):
         if isinstance(self.server, (list, tuple)):
             sock = self.socket_module.socket(self.socket_module.AF_INET,
                                              self.socket_module.SOCK_STREAM)
+
+            if self.tls_context:
+                sock = self.tls_context.wrap_socket(
+                    sock, server_hostname=self.server[0]
+                )
         else:
             sock = self.socket_module.socket(self.socket_module.AF_UNIX,
                                              self.socket_module.SOCK_STREAM)
@@ -291,6 +298,7 @@ class Client(object):
             if self.no_delay and sock.family == self.socket_module.AF_INET:
                 sock.setsockopt(self.socket_module.IPPROTO_TCP,
                                 self.socket_module.TCP_NODELAY, 1)
+
         except Exception:
             sock.close()
             raise
@@ -1016,7 +1024,8 @@ class PooledClient(object):
                  lock_generator=None,
                  default_noreply=True,
                  allow_unicode_keys=False,
-                 encoding='ascii'):
+                 encoding='ascii',
+                 tls_context=None):
         self.server = server
         self.serde = serde or LegacyWrappingSerde(serializer, deserializer)
         self.connect_timeout = connect_timeout
@@ -1037,6 +1046,7 @@ class PooledClient(object):
             max_size=max_pool_size,
             lock_generator=lock_generator)
         self.encoding = encoding
+        self.tls_context = tls_context
 
     def check_key(self, key):
         """Checks key and add key_prefix."""
@@ -1055,7 +1065,8 @@ class PooledClient(object):
                         socket_module=self.socket_module,
                         key_prefix=self.key_prefix,
                         default_noreply=self.default_noreply,
-                        allow_unicode_keys=self.allow_unicode_keys)
+                        allow_unicode_keys=self.allow_unicode_keys,
+                        tls_context=self.tls_context)
         return client
 
     def close(self):
diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py
index fccb060..25e297e 100644
--- a/pymemcache/client/hash.py
+++ b/pymemcache/client/hash.py
@@ -36,7 +36,8 @@ class HashClient(object):
         ignore_exc=False,
         allow_unicode_keys=False,
         default_noreply=True,
-        encoding='ascii'
+        encoding='ascii',
+        tls_context=None
     ):
         """
         Constructor.
@@ -87,7 +88,8 @@ class HashClient(object):
             'deserializer': deserializer,
             'allow_unicode_keys': allow_unicode_keys,
             'default_noreply': default_noreply,
-            'encoding': encoding
+            'encoding': encoding,
+            'tls_context': tls_context,
         }
 
         if use_pooling is True:
@@ -99,6 +101,7 @@ class HashClient(object):
         for server, port in servers:
             self.add_server(server, port)
         self.encoding = encoding
+        self.tls_context = tls_context
 
     def add_server(self, server, port):
         key = '%s:%s' % (server, port)