Merge pull request #118 from pinterest/fix-unicode-char-in-middle-error

Fix #117, illegal unicode character in middle of key (or value now)

.gitignore

@@ -40,6 +40,7 @@ env/
 .pip
 .pypirc
 coverage.xml
+.python-version
 \#*\#
 docs/_build
 docs/apidoc/pymemcache.test.rst

pymemcache/client/base.py

@@ -82,10 +82,11 @@ def _parse_hex(value):
 
 def _check_key(key, key_prefix=b''):
     """Checks key and add key_prefix."""
-    if isinstance(key, six.text_type):
+    allowed_str_types = (six.text_type, six.string_types)
+    if isinstance(key, allowed_str_types):
         try:
             key = key.encode('ascii')
-        except UnicodeEncodeError:
+        except (UnicodeEncodeError, UnicodeDecodeError):
             raise MemcacheIllegalInputError("Non-ASCII key: '%r'" % (key,))
     key = key_prefix + key
     if b' ' in key or b'\n' in key:

pymemcache/test/test_client.py

@@ -92,6 +92,15 @@ def _set():
         with pytest.raises(MemcacheIllegalInputError):
             _set()
 
+    def test_set_unicode_char_in_middle_of_key(self):
+        client = self.make_client([b'STORED\r\n'])
+
+        def _set():
+            client.set('helloworld_\xb1901520_%c3', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
     def test_set_unicode_value(self):
         client = self.make_client([b''])
 

pymemcache/test/utils.py

@@ -43,6 +43,11 @@ def __init__(self,
     def get(self, key, default=None):
         if isinstance(key, six.text_type):
             raise MemcacheIllegalInputError(key)
+        if isinstance(key, six.string_types):
+            try:
+                key = key.encode('ascii')
+            except (UnicodeEncodeError, UnicodeDecodeError):
+                raise MemcacheIllegalInputError
 
         if key not in self._contents:
             return default
@@ -71,6 +76,16 @@ def set(self, key, value, expire=0, noreply=True):
             raise MemcacheIllegalInputError(key)
         if isinstance(value, six.text_type):
             raise MemcacheIllegalInputError(value)
+        if isinstance(key, six.string_types):
+            try:
+                key = key.encode('ascii')
+            except (UnicodeEncodeError, UnicodeDecodeError):
+                raise MemcacheIllegalInputError
+        if isinstance(value, six.string_types):
+            try:
+                value = value.encode('ascii')
+            except (UnicodeEncodeError, UnicodeDecodeError):
+                raise MemcacheIllegalInputError
 
         flags = 0
         if self.serializer:

tox.ini

@@ -1,5 +1,6 @@
 [tox]
-envlist = py26, py27, pypy, pypy3, py33, py34, docs, py27-flake8, py34-flake8
+envlist = py26, py27, pypy, pypy3, py33, py34, py35, docs, py27-flake8, py35-flake8
+skip_missing_interpreters = True
 
 [testenv]
 commands =
@@ -12,7 +13,7 @@ commands =
     pip install flake8
     flake8 pymemcache/
 
-[testenv:py34-flake8]
+[testenv:py35-flake8]
 commands =
     pip install flake8
     flake8 pymemcache/