Skip to content

Commit

Permalink
Add optional client_name property to RedisConnectionInfo, which will …
Browse files Browse the repository at this point in the history
…be used with 'CLIENT SETNAME' command during connection setup.
  • Loading branch information
ikolomi committed Dec 11, 2023
1 parent 46e8a9d commit 0cf251e
Show file tree
Hide file tree
Showing 6 changed files with 182 additions and 44 deletions.
15 changes: 15 additions & 0 deletions redis/src/aio/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
23 changes: 23 additions & 0 deletions redis/src/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,8 @@ pub struct RedisConnectionInfo {
pub password: Option<String>,
/// 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<String>,
}

impl FromStr for ConnectionInfo {
Expand Down Expand Up @@ -387,6 +389,7 @@ fn url_to_tcp_connection_info(url: url::Url) -> RedisResult<ConnectionInfo> {
Some(v) => v == "true",
_ => false,
},
client_name: None,
},
})
}
Expand All @@ -413,6 +416,7 @@ fn url_to_unix_connection_info(url: url::Url) -> RedisResult<ConnectionInfo> {
Some(v) => v == "true",
_ => false,
},
client_name: None,
},
})
}
Expand Down Expand Up @@ -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::<Value>(&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);
Expand Down Expand Up @@ -1708,6 +1726,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("#@<>$".to_string()),
use_resp3: false,
client_name: None,
},
},
),
Expand Down Expand Up @@ -1775,6 +1794,7 @@ mod tests {
username: None,
password: None,
use_resp3: false,
client_name: None,
},
},
),
Expand All @@ -1787,6 +1807,7 @@ mod tests {
username: None,
password: None,
use_resp3: false,
client_name: None,
},
},
),
Expand All @@ -1802,6 +1823,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("#@<>$".to_string()),
use_resp3: false,
client_name: None,
},
},
),
Expand All @@ -1817,6 +1839,7 @@ mod tests {
username: Some("%johndoe%".to_string()),
password: Some("&?= *+".to_string()),
use_resp3: false,
client_name: None,
},
},
),
Expand Down
2 changes: 2 additions & 0 deletions redis/src/sentinel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
//! username: Some(String::from("foo")),
//! password: Some(String::from("bar")),
//! use_resp3: false,
//! client_name: None
//! }),
//! }),
//! )
Expand Down Expand Up @@ -95,6 +96,7 @@
//! username: Some(String::from("user")),
//! password: Some(String::from("pass")),
//! use_resp3: false,
//! client_name: None
//! }),
//! }),
//! redis::sentinel::SentinelServerType::Master,
Expand Down
96 changes: 52 additions & 44 deletions redis/tests/support/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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,
Expand All @@ -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,
Expand Down
50 changes: 50 additions & 0 deletions redis/tests/test_async.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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();
}
40 changes: 40 additions & 0 deletions redis/tests/test_basic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
);
}

0 comments on commit 0cf251e

Please sign in to comment.