Raising MemcacheIllegalInputError when a key contains null byte, new line, or carriage return

2 files changed, 36 insertions, 1 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 60b9f2a..a842455 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -96,6 +96,14 @@ def _check_key(key, allow_unicode_keys, key_prefix=b''):
         raise MemcacheIllegalInputError(
           "Key contains space and/or newline: '%r'" % (key,)
         )
+    if b'\00' in key:
+        raise MemcacheIllegalInputError(
+          "Key contains null byte: '%r'" % (key,)
+        )
+    if b'\r\n' in key:
+        raise MemcacheIllegalInputError(
+          "Key contains carriage return: '%r'" % (key,)
+        )
     if len(key) > 250:
         raise MemcacheIllegalInputError("Key is too long: '%r'" % (key,))
     return key
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index ea7b044..05e29a0 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -132,6 +132,33 @@ class ClientTestMixin(object):
         with pytest.raises(MemcacheIllegalInputError):
             _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\n', b'value', noreply=False)
+
+        with pytest.raises(MemcacheIllegalInputError):
+            _set()
+
     def test_set_unicode_char_in_middle_of_key_ok(self):
         client = self.make_client([b'STORED\r\n'], allow_unicode_keys=True)
 
@@ -664,7 +691,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')