Iterating over characters in key once

2 files changed, 20 insertions, 13 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index a842455..7b365f0 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -92,20 +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 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,))
+
+    for c in key:
+        if c == b' ':
+            raise MemcacheIllegalInputError(
+                "Key contains space: '%r'" % (key,)
+            )
+        elif c == b'\n':
+            raise MemcacheIllegalInputError(
+                "Key contains newline: '%r'" % (key,)
+            )
+        elif c == b'\00':
+            raise MemcacheIllegalInputError(
+              "Key contains null byte: '%r'" % (key,)
+            )
+        elif c == 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 329562c..8d17411 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -576,7 +576,7 @@ class TestClient(ClientTestMixin, unittest.TestCase):
         client = self.make_client([b''])
 
         def _set():
-            client.set(b'key\r\n', b'value', noreply=False)
+            client.set(b'key\r', b'value', noreply=False)
 
         with pytest.raises(MemcacheIllegalInputError):
             _set()