Merge pull request #138 from bwalks/master

Raising MemcacheIllegalInputError when key contains null byte, newline, or carriage return

2 files changed, 56 insertions, 5 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 60b9f2a..3578058 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -92,12 +92,27 @@ def _check_key(key, allow_unicode_keys, key_prefix=b''):
         except (UnicodeEncodeError, UnicodeDecodeError):
             raise MemcacheIllegalInputError("Non-ASCII key: '%r'" % (key,))
     key = key_prefix + key
-    if b' ' in key or b'\n' in key:
-        raise MemcacheIllegalInputError(
-          "Key contains space and/or newline: '%r'" % (key,)
-        )
+
     if len(key) > 250:
         raise MemcacheIllegalInputError("Key is too long: '%r'" % (key,))
+
+    for c in bytearray(key):
+        if c == ord(b' '):
+            raise MemcacheIllegalInputError(
+                "Key contains space: '%r'" % (key,)
+            )
+        elif c == ord(b'\n'):
+            raise MemcacheIllegalInputError(
+                "Key contains newline: '%r'" % (key,)
+            )
+        elif c == ord(b'\00'):
+            raise MemcacheIllegalInputError(
+              "Key contains null character: '%r'" % (key,)
+            )
+        elif c == ord(b'\r'):
+            raise MemcacheIllegalInputError(
+              "Key contains carriage return: '%r'" % (key,)
+            )
     return key
 
 
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index ea7b044..606a950 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -554,6 +554,42 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         with pytest.raises(MemcacheUnknownError):
             _set()
 
+    def test_set_key_with_space(self):
+        client = self.make_client([b''])
+
+        def _set():
+            client.set(b'key has space', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
+    def test_set_key_with_newline(self):
+        client = self.make_client([b''])
+
+        def _set():
+            client.set(b'key\n', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
+    def test_set_key_with_carriage_return(self):
+        client = self.make_client([b''])
+
+        def _set():
+            client.set(b'key\r', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
+    def test_set_key_with_null_character(self):
+        client = self.make_client([b''])
+
+        def _set():
+            client.set(b'key\00', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
     def test_set_many_socket_handling(self):
         client = self.make_client([b'STORED\r\n'])
         result = client.set_many({b'key': b'value'}, noreply=False)
@@ -664,7 +700,7 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         with pytest.raises(MemcacheClientError):
             client.get(u'\u0FFF'*150)
 
-    def test_key_contains_spae(self):
+    def test_key_contains_space(self):
         client = self.make_client([b'END\r\n'])
         with pytest.raises(MemcacheClientError):
             client.get(b'abc xyz')