Repository: aisk/rust-memcache (language: rust)

Issue:
Move check_key_len to Client top-level.
Memcached's max key len is 256 bytes. We check keys' length before setting them in the server. Right now, it is done just before sending the payload to the server. Instead, we should check keys for this as soon as we get input and fail fast.

Example: https://github.com/aisk/rust-memcache/blob/faed4e117b8108dbcae1901c891844009229190e/src/protocol/ascii.rs#L151


Discussion / hints:

Fix patch:
diff --git a/src/client.rs b/src/client.rs
index 8cd5ab0..14c9adc 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -6,7 +6,7 @@ use std::time::Duration;
 use url::Url;
 
 use connection::ConnectionManager;
-use error::MemcacheError;
+use error::{ClientError, MemcacheError};
 use protocol::{Protocol, ProtocolTrait};
 use r2d2::Pool;
 use stream::Stream;
@@ -60,6 +60,13 @@ fn default_hash_function(key: &str) -> u64 {
     return hasher.finish();
 }
 
+pub(crate) fn check_key_len(key: &str) -> Result<(), MemcacheError> {
+    if key.len() > 250 {
+        Err(ClientError::KeyTooLong)?
+    }
+    Ok(())
+}
+
 impl Client {
     #[deprecated(since = "0.10.0", note = "please use `connect` instead")]
     pub fn new<C: Connectable>(target: C) -> Result<Self, MemcacheError> {
@@ -186,6 +193,7 @@ impl Client {
     /// let _: Option<String> = client.get("foo").unwrap();
     /// ```
     pub fn get<V: FromMemcacheValueExt>(&self, key: &str) -> Result<Option<V>, MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.get(key);
     }
 
@@ -201,6 +209,9 @@ impl Client {
     /// assert_eq!(result["foo"], "42");
     /// ```
     pub fn gets<V: FromMemcacheValueExt>(&self, keys: &[&str]) -> Result<HashMap<String, V>, MemcacheError> {
+        for key in keys {
+            check_key_len(key)?;
+        }
         let mut con_keys: HashMap<usize, Vec<&str>> = HashMap::new();
         let mut result: HashMap<String, V> = HashMap::new();
         let connections_count = self.connections.len();
@@ -227,6 +238,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn set<V: ToMemcacheValue<Stream>>(&self, key: &str, value: V, expiration: u32) -> Result<(), MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.set(key, value, expiration);
     }
 
@@ -252,6 +264,7 @@ impl Client {
         expiration: u32,
         cas_id: u64,
     ) -> Result<bool, MemcacheError> {
+        check_key_len(key)?;
         self.get_connection(key).get()?.cas(key, value, expiration, cas_id)
     }
 
@@ -267,6 +280,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn add<V: ToMemcacheValue<Stream>>(&self, key: &str, value: V, expiration: u32) -> Result<(), MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.add(key, value, expiration);
     }
 
@@ -287,6 +301,7 @@ impl Client {
         value: V,
         expiration: u32,
     ) -> Result<(), MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.replace(key, value, expiration);
     }
 
@@ -304,6 +319,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn append<V: ToMemcacheValue<Stream>>(&self, key: &str, value: V) -> Result<(), MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.append(key, value);
     }
 
@@ -321,6 +337,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn prepend<V: ToMemcacheValue<Stream>>(&self, key: &str, value: V) -> Result<(), MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.prepend(key, value);
     }
 
@@ -334,6 +351,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn delete(&self, key: &str) -> Result<bool, MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.delete(key);
     }
 
@@ -347,6 +365,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn increment(&self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.increment(key, amount);
     }
 
@@ -360,6 +379,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn decrement(&self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.decrement(key, amount);
     }
 
@@ -375,6 +395,7 @@ impl Client {
     /// # client.flush().unwrap();
     /// ```
     pub fn touch(&self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
+        check_key_len(key)?;
         return self.get_connection(key).get()?.touch(key, expiration);
     }
 
diff --git a/src/protocol/ascii.rs b/src/protocol/ascii.rs
index 45014b3..346c11a 100644
--- a/src/protocol/ascii.rs
+++ b/src/protocol/ascii.rs
@@ -2,7 +2,6 @@ use std::collections::HashMap;
 use std::fmt;
 use std::io::{Read, Write};
 
-use super::check_key_len;
 use client::Stats;
 use error::{ClientError, CommandError, MemcacheError, ServerError};
 use std::borrow::Cow;
@@ -148,7 +147,6 @@ impl AsciiProtocol<Stream> {
         value: V,
         options: &Options,
     ) -> Result<bool, MemcacheError> {
-        check_key_len(key)?;
         if command == StoreCommand::Cas {
             if options.cas.is_none() {
                 Err(ClientError::Error(Cow::Borrowed(
@@ -302,9 +300,6 @@ impl AsciiProtocol<Stream> {
     }
 
     pub(super) fn gets<V: FromMemcacheValueExt>(&mut self, keys: &[&str]) -> Result<HashMap<String, V>, MemcacheError> {
-        for key in keys {
-            check_key_len(key)?;
-        }
         write!(self.reader.get_mut(), "gets {}\r\n", keys.join(" "))?;
 
         let mut result: HashMap<String, V> = HashMap::with_capacity(keys.len());
@@ -382,19 +377,16 @@ impl AsciiProtocol<Stream> {
     }
 
     pub(super) fn append<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
-        check_key_len(key)?;
         self.store(StoreCommand::Append, key, value, &Default::default())
             .map(|_| ())
     }
 
     pub(super) fn prepend<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
-        check_key_len(key)?;
         self.store(StoreCommand::Prepend, key, value, &Default::default())
             .map(|_| ())
     }
 
     pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
-        check_key_len(key)?;
         write!(self.reader.get_mut(), "delete {}\r\n", key)?;
         self.reader.get_mut().flush()?;
         self.reader
@@ -419,19 +411,16 @@ impl AsciiProtocol<Stream> {
     }
 
     pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
-        check_key_len(key)?;
         write!(self.reader.get_mut(), "incr {} {}\r\n", key, amount)?;
         self.parse_u64_response()
     }
 
     pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
-        check_key_len(key)?;
         write!(self.reader.get_mut(), "decr {} {}\r\n", key, amount)?;
         self.parse_u64_response()
     }
 
     pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
-        check_key_len(key)?;
         write!(self.reader.get_mut(), "touch {} {}\r\n", key, expiration)?;
         self.reader.get_mut().flush()?;
         self.reader
diff --git a/src/protocol/binary.rs b/src/protocol/binary.rs
index 39367f8..dc8d74c 100644
--- a/src/protocol/binary.rs
+++ b/src/protocol/binary.rs
@@ -1,7 +1,6 @@
 use std::collections::HashMap;
 use std::io::Write;
 
-use super::check_key_len;
 use byteorder::{BigEndian, WriteBytesExt};
 use client::Stats;
 use error::MemcacheError;
@@ -38,7 +37,6 @@ impl BinaryProtocol {
         expiration: u32,
         cas: Option<u64>,
     ) -> Result<(), MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: opcode as u8,
@@ -110,7 +108,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn get<V: FromMemcacheValueExt>(&mut self, key: &str) -> Result<Option<V>, MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Get as u8,
@@ -126,7 +123,6 @@ impl BinaryProtocol {
 
     pub(super) fn gets<V: FromMemcacheValueExt>(&mut self, keys: &[&str]) -> Result<HashMap<String, V>, MemcacheError> {
         for key in keys {
-            check_key_len(key)?;
             let request_header = PacketHeader {
                 magic: Magic::Request as u8,
                 opcode: Opcode::GetKQ as u8,
@@ -185,7 +181,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn append<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Append as u8,
@@ -201,7 +196,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn prepend<V: ToMemcacheValue<Stream>>(&mut self, key: &str, value: V) -> Result<(), MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Prepend as u8,
@@ -217,7 +211,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn delete(&mut self, key: &str) -> Result<bool, MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Delete as u8,
@@ -232,7 +225,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn increment(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Increment as u8,
@@ -256,7 +248,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn decrement(&mut self, key: &str, amount: u64) -> Result<u64, MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Decrement as u8,
@@ -280,7 +271,6 @@ impl BinaryProtocol {
     }
 
     pub(super) fn touch(&mut self, key: &str, expiration: u32) -> Result<bool, MemcacheError> {
-        check_key_len(key)?;
         let request_header = PacketHeader {
             magic: Magic::Request as u8,
             opcode: Opcode::Touch as u8,
diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 56d667b..77e5664 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -2,7 +2,6 @@ mod ascii;
 mod binary;
 mod binary_packet;
 
-use crate::error::ClientError;
 use client::Stats;
 use enum_dispatch::enum_dispatch;
 use error::MemcacheError;
@@ -12,13 +11,6 @@ use std::collections::HashMap;
 use stream::Stream;
 use value::{FromMemcacheValueExt, ToMemcacheValue};
 
-pub(crate) fn check_key_len(key: &str) -> Result<(), MemcacheError> {
-    if key.len() > 250 {
-        Err(ClientError::KeyTooLong)?
-    }
-    Ok(())
-}
-
 #[enum_dispatch]
 pub enum Protocol {
     Ascii(AsciiProtocol<Stream>),


Test patch: