diff --git a/src/client.rs b/src/client.rs index c4e5957..5ba24a3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -78,10 +78,13 @@ fn default_hash_function(key: &str) -> u64 { return hasher.finish(); } -pub(crate) fn check_key_len(key: &str) -> Result<(), MemcacheError> { +pub(crate) fn check_key(key: &str) -> Result<(), MemcacheError> { if key.len() > 250 { Err(ClientError::KeyTooLong)? } + if key.bytes().any(|b| b <= b' ' || b == 0x7f) { + Err(ClientError::InvalidKey)? + } Ok(()) } @@ -237,7 +240,7 @@ impl Client { /// let _: Option = client.get("foo").unwrap(); /// ``` pub fn get(&self, key: &str) -> Result, MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.get(key)); } @@ -254,7 +257,7 @@ impl Client { /// ``` pub fn gets(&self, keys: &[&str]) -> Result, MemcacheError> { for key in keys { - check_key_len(key)?; + check_key(key)?; } let mut con_keys: HashMap> = HashMap::new(); let mut result: HashMap = HashMap::new(); @@ -281,7 +284,7 @@ impl Client { /// client.flush().unwrap(); /// ``` pub fn set>(&self, key: &str, value: V, expiration: u32) -> Result<(), MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.set(key, value, expiration)); } @@ -305,7 +308,7 @@ impl Client { expiration: u32, cas_id: u64, ) -> Result { - check_key_len(key)?; + check_key(key)?; with_connection(&self.get_connection(key), |c| c.cas(key, value, expiration, cas_id)) } @@ -321,7 +324,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn add>(&self, key: &str, value: V, expiration: u32) -> Result<(), MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.add(key, value, expiration)); } @@ -342,7 +345,7 @@ impl Client { value: V, expiration: u32, ) -> Result<(), MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.replace(key, value, expiration)); } @@ -360,7 +363,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn append>(&self, key: &str, value: V) -> Result<(), MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.append(key, value)); } @@ -378,7 +381,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn prepend>(&self, key: &str, value: V) -> Result<(), MemcacheError> { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.prepend(key, value)); } @@ -392,7 +395,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn delete(&self, key: &str) -> Result { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.delete(key)); } @@ -406,7 +409,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn increment(&self, key: &str, amount: u64) -> Result { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.increment(key, amount)); } @@ -420,7 +423,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn decrement(&self, key: &str, amount: u64) -> Result { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.decrement(key, amount)); } @@ -436,7 +439,7 @@ impl Client { /// # client.flush().unwrap(); /// ``` pub fn touch(&self, key: &str, expiration: u32) -> Result { - check_key_len(key)?; + check_key(key)?; return with_connection(&self.get_connection(key), |c| c.touch(key, expiration)); } @@ -594,6 +597,28 @@ impl ClientBuilder { mod tests { use std::time::Duration; + #[test] + fn check_key() { + use crate::error::{ClientError, MemcacheError}; + + assert!(super::check_key("foo").is_ok()); + assert!(super::check_key(&"k".repeat(250)).is_ok()); + assert!(matches!( + super::check_key(&"k".repeat(251)), + Err(MemcacheError::ClientError(ClientError::KeyTooLong)) + )); + for key in ["foo bar", "foo\r\nflush_all", "foo\n", "\tfoo", "foo\0", "foo\x7f"] { + assert!( + matches!( + super::check_key(key), + Err(MemcacheError::ClientError(ClientError::InvalidKey)) + ), + "{:?}", + key + ); + } + } + #[test] fn build_client_happy_path() { let client = super::Client::builder() diff --git a/src/error.rs b/src/error.rs index ab087d9..aab511a 100644 --- a/src/error.rs +++ b/src/error.rs @@ -11,6 +11,8 @@ use std::string; pub enum ClientError { /// The key provided was longer than 250 bytes. KeyTooLong, + /// The key provided contained whitespace or control characters. + InvalidKey, /// The server returned an error prefixed with CLIENT_ERROR in response to a command. Error(Cow<'static, str>), } @@ -19,6 +21,7 @@ impl fmt::Display for ClientError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { ClientError::KeyTooLong => write!(f, "The provided key was too long."), + ClientError::InvalidKey => write!(f, "The provided key contained whitespace or control characters."), ClientError::Error(s) => write!(f, "{}", s), } } diff --git a/tests/test_ascii.rs b/tests/test_ascii.rs index 38ad7a6..ff5573f 100644 --- a/tests/test_ascii.rs +++ b/tests/test_ascii.rs @@ -15,6 +15,11 @@ fn test_ascii() { let value: Option = client.get("ascii_foo").unwrap(); assert_eq!(value, Some("bar".into())); + assert!(client.get::("ascii_foo\r\nflush_all").is_err()); + assert!(client.set("ascii foo", "bar", 0).is_err()); + let value: Option = client.get("ascii_foo").unwrap(); + assert_eq!(value, Some("bar".into())); + client.set("ascii_baz", "qux", 0).unwrap(); let values: HashMap, u32)> = client.gets(&["ascii_foo", "ascii_baz", "not_exists_key"]).unwrap(); assert_eq!(values.len(), 2);