From 0cf251e613ba8c2da11dfa7b9d81a31dd3d25358 Mon Sep 17 00:00:00 2001 From: ikolomi Date: Sun, 10 Dec 2023 13:52:15 +0200 Subject: [PATCH] Add optional client_name property to RedisConnectionInfo, which will be used with 'CLIENT SETNAME' command during connection setup. --- redis/src/aio/mod.rs | 15 ++++++ redis/src/connection.rs | 23 +++++++++ redis/src/sentinel.rs | 2 + redis/tests/support/mod.rs | 96 +++++++++++++++++++++----------------- redis/tests/test_async.rs | 50 ++++++++++++++++++++ redis/tests/test_basic.rs | 40 ++++++++++++++++ 6 files changed, 182 insertions(+), 44 deletions(-) diff --git a/redis/src/aio/mod.rs b/redis/src/aio/mod.rs index b64f06d3b6..24347c2fb4 100644 --- a/redis/src/aio/mod.rs +++ b/redis/src/aio/mod.rs @@ -140,6 +140,21 @@ where } } + if let Some(client_name) = &connection_info.client_name { + match cmd("CLIENT") + .arg("SETNAME") + .arg(client_name) + .query_async(con) + .await + { + Ok(Value::Okay) => {} + _ => fail!(( + ErrorKind::ResponseError, + "Redis server refused to set client name" + )), + } + } + // result is ignored, as per the command's instructions. // https://redis.io/commands/client-setinfo/ let _: RedisResult<()> = crate::connection::client_set_info_pipeline() diff --git a/redis/src/connection.rs b/redis/src/connection.rs index 92e249dfa1..f3ac856633 100644 --- a/redis/src/connection.rs +++ b/redis/src/connection.rs @@ -227,6 +227,8 @@ pub struct RedisConnectionInfo { pub password: Option, /// Use RESP 3 mode, Redis 6 or newer is required. pub use_resp3: bool, + /// Optionally a pass a client name that should be used for connection + pub client_name: Option, } impl FromStr for ConnectionInfo { @@ -387,6 +389,7 @@ fn url_to_tcp_connection_info(url: url::Url) -> RedisResult { Some(v) => v == "true", _ => false, }, + client_name: None, }, }) } @@ -413,6 +416,7 @@ fn url_to_unix_connection_info(url: url::Url) -> RedisResult { Some(v) => v == "true", _ => false, }, + client_name: None, }, }) } @@ -979,6 +983,20 @@ fn setup_connection( } } + if connection_info.client_name.is_some() { + match cmd("CLIENT") + .arg("SETNAME") + .arg(connection_info.client_name.as_ref().unwrap()) + .query::(&mut rv) + { + Ok(Value::Okay) => {} + _ => fail!(( + ErrorKind::ResponseError, + "Redis server refused to set client name" + )), + } + } + // result is ignored, as per the command's instructions. // https://redis.io/commands/client-setinfo/ let _: RedisResult<()> = client_set_info_pipeline().query(&mut rv); @@ -1708,6 +1726,7 @@ mod tests { username: Some("%johndoe%".to_string()), password: Some("#@<>$".to_string()), use_resp3: false, + client_name: None, }, }, ), @@ -1775,6 +1794,7 @@ mod tests { username: None, password: None, use_resp3: false, + client_name: None, }, }, ), @@ -1787,6 +1807,7 @@ mod tests { username: None, password: None, use_resp3: false, + client_name: None, }, }, ), @@ -1802,6 +1823,7 @@ mod tests { username: Some("%johndoe%".to_string()), password: Some("#@<>$".to_string()), use_resp3: false, + client_name: None, }, }, ), @@ -1817,6 +1839,7 @@ mod tests { username: Some("%johndoe%".to_string()), password: Some("&?= *+".to_string()), use_resp3: false, + client_name: None, }, }, ), diff --git a/redis/src/sentinel.rs b/redis/src/sentinel.rs index 9a32ac8898..2045894fb9 100644 --- a/redis/src/sentinel.rs +++ b/redis/src/sentinel.rs @@ -59,6 +59,7 @@ //! username: Some(String::from("foo")), //! password: Some(String::from("bar")), //! use_resp3: false, +//! client_name: None //! }), //! }), //! ) @@ -95,6 +96,7 @@ //! username: Some(String::from("user")), //! password: Some(String::from("pass")), //! use_resp3: false, +//! client_name: None //! }), //! }), //! redis::sentinel::SentinelServerType::Master, diff --git a/redis/tests/support/mod.rs b/redis/tests/support/mod.rs index 9c2fdb5fbf..1cef28f43f 100644 --- a/redis/tests/support/mod.rs +++ b/redis/tests/support/mod.rs @@ -349,28 +349,7 @@ impl TestContext { Self::with_modules(&[], true) } - pub fn with_tls(tls_files: TlsFilePaths, mtls_enabled: bool) -> TestContext { - let redis_port = get_random_available_port(); - let addr = RedisServer::get_addr(redis_port); - - let server = RedisServer::new_with_addr_tls_modules_and_spawner( - addr, - None, - Some(tls_files), - mtls_enabled, - &[], - |cmd| { - cmd.spawn() - .unwrap_or_else(|err| panic!("Failed to run {cmd:?}: {err}")) - }, - ); - - #[cfg(feature = "tls-rustls")] - let client = - build_single_client(server.connection_info(), &server.tls_paths, mtls_enabled).unwrap(); - #[cfg(not(feature = "tls-rustls"))] - let client = redis::Client::open(server.connection_info()).unwrap(); - + fn connect_with_retries(client: &redis::Client) { let mut con; let millisecond = Duration::from_millis(1); @@ -395,6 +374,31 @@ impl TestContext { } } redis::cmd("FLUSHDB").execute(&mut con); + } + + pub fn with_tls(tls_files: TlsFilePaths, mtls_enabled: bool) -> TestContext { + let redis_port = get_random_available_port(); + let addr: ConnectionAddr = RedisServer::get_addr(redis_port); + + let server = RedisServer::new_with_addr_tls_modules_and_spawner( + addr, + None, + Some(tls_files), + mtls_enabled, + &[], + |cmd| { + cmd.spawn() + .unwrap_or_else(|err| panic!("Failed to run {cmd:?}: {err}")) + }, + ); + + #[cfg(feature = "tls-rustls")] + let client = + build_single_client(server.connection_info(), &server.tls_paths, mtls_enabled).unwrap(); + #[cfg(not(feature = "tls-rustls"))] + let client = redis::Client::open(server.connection_info()).unwrap(); + + Self::connect_with_retries(&client); TestContext { server, @@ -412,30 +416,34 @@ impl TestContext { #[cfg(not(feature = "tls-rustls"))] let client = redis::Client::open(server.connection_info()).unwrap(); - let mut con; + Self::connect_with_retries(&client); - let millisecond = Duration::from_millis(1); - let mut retries = 0; - loop { - match client.get_connection() { - Err(err) => { - if err.is_connection_refusal() { - sleep(millisecond); - retries += 1; - if retries > 100000 { - panic!("Tried to connect too many times, last error: {err}"); - } - } else { - panic!("Could not connect: {err}"); - } - } - Ok(x) => { - con = x; - break; - } - } + TestContext { + server, + client, + use_resp3: use_resp3(), } - redis::cmd("FLUSHDB").execute(&mut con); + } + + pub fn with_client_name(clientname: &str) -> TestContext { + let server = RedisServer::with_modules(&[], false); + let con_info = redis::ConnectionInfo { + addr: server.client_addr().clone(), + redis: redis::RedisConnectionInfo { + db: Default::default(), + username: None, + password: None, + use_resp3: Default::default(), + client_name: Some(clientname.to_string()), + }, + }; + + #[cfg(feature = "tls-rustls")] + let client = build_single_client(con_info, &server.tls_paths, false).unwrap(); + #[cfg(not(feature = "tls-rustls"))] + let client = redis::Client::open(con_info).unwrap(); + + Self::connect_with_retries(&client); TestContext { server, diff --git a/redis/tests/test_async.rs b/redis/tests/test_async.rs index d9ca2b7642..ba66b87595 100644 --- a/redis/tests/test_async.rs +++ b/redis/tests/test_async.rs @@ -472,6 +472,7 @@ async fn invalid_password_issue_343() { username: None, password: Some("asdcasc".to_string()), use_resp3: false, + client_name: None, }, }; let client = redis::Client::open(coninfo).unwrap(); @@ -781,3 +782,52 @@ mod mtls_test { } } } + +#[test] +fn test_set_client_name_by_config() { + const CLIENT_NAME: &str = "TEST_CLIENT_NAME"; + use redis::RedisError; + let ctx = TestContext::with_client_name(CLIENT_NAME); + + block_on_all(async move { + let mut con = ctx.async_connection().await?; + + let client_list: redis::InfoDict = redis::cmd("CLIENT") + .arg("LIST") + .query_async(&mut con) + .await + .unwrap(); + assert!( + !client_list.is_empty(), + "No clients present in CLIENT LIST result" + ); + + let val: &redis::Value = client_list.values().next().unwrap(); + + match val { + redis::Value::SimpleString(s) => { + let client_attr: Vec<&str> = s.split(' ').collect(); + for i in client_attr.iter() { + let this_attr: Vec<&str> = i.split('=').collect(); + if this_attr[0] == "name" { + assert!( + this_attr[1] == CLIENT_NAME, + "Incorrect client name, expecting: {}, got {}", + CLIENT_NAME, + this_attr[1] + ); + return Ok::<_, RedisError>(()); + } + } + } + _ => { + panic!("Unexpected Enum type returned by CLIENT LIST command (expected String)"); + } + } + panic!( + "Could not detect the expected client name {} in CLIENT LIST output", + CLIENT_NAME + ); + }) + .unwrap(); +} diff --git a/redis/tests/test_basic.rs b/redis/tests/test_basic.rs index 48da016985..554667b6ef 100644 --- a/redis/tests/test_basic.rs +++ b/redis/tests/test_basic.rs @@ -1448,3 +1448,43 @@ fn test_blocking_sorted_set_api() { ); } } + +#[test] +fn test_set_client_name_by_config() { + const CLIENT_NAME: &str = "TEST_CLIENT_NAME"; + + let ctx = TestContext::with_client_name(CLIENT_NAME); + let mut con = ctx.connection(); + + let client_list: redis::InfoDict = redis::cmd("CLIENT").arg("LIST").query(&mut con).unwrap(); + assert!( + !client_list.is_empty(), + "No clients present in CLIENT LIST result" + ); + + let val: &redis::Value = client_list.values().next().unwrap(); + match val { + redis::Value::SimpleString(s) => { + let client_attr: Vec<&str> = s.split(' ').collect(); + for i in client_attr.iter() { + let this_attr: Vec<&str> = i.split('=').collect(); + if this_attr[0] == "name" { + assert!( + this_attr[1] == CLIENT_NAME, + "Incorrect client name, expecting: {}, got {}", + CLIENT_NAME, + this_attr[1] + ); + return; + } + } + } + _ => { + panic!("Unexpected Enum type returned by CLIENT LIST command (expected String)"); + } + } + panic!( + "Could not detect the expected client name {} in CLIENT LIST output", + CLIENT_NAME + ); +}