From e354aca3c1bf4f055242a1ff8303a446277f335c Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Wed, 25 Feb 2026 11:36:41 -0500 Subject: [PATCH] refactor(cli): replace duplicated connection code with ember-client MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ember-cli had its own connection.rs with ~85% overlap with ember-client/src/connection.rs — identical error types, frame read loop, send/auth/disconnect logic, and TlsClientConfig struct. - expose TlsClientConfig publicly from ember-client (tls feature) - re-export it from ember-cli::tls, removing the duplicate struct - add ember-client dependency to ember-cli - update all callers (repl, batch, watch, cluster, main) to use Client::connect / Client::connect_tls, .send(), .auth(), .disconnect() - inline auth_frame in bench_conn.rs (4-line helper, one call site) - delete crates/ember-cli/src/connection.rs --- Cargo.lock | 114 +++++++++++++++-- Cargo.toml | 1 + crates/ember-cli/Cargo.toml | 1 + crates/ember-cli/src/batch.rs | 21 ++-- crates/ember-cli/src/bench_conn.rs | 6 +- crates/ember-cli/src/cluster.rs | 96 +++++++------- crates/ember-cli/src/connection.rs | 193 ----------------------------- crates/ember-cli/src/main.rs | 16 ++- crates/ember-cli/src/repl.rs | 32 +++-- crates/ember-cli/src/tls.rs | 18 +-- crates/ember-cli/src/watch.rs | 16 ++- crates/ember-client/src/lib.rs | 2 + 12 files changed, 225 insertions(+), 291 deletions(-) delete mode 100644 crates/ember-cli/src/connection.rs diff --git a/Cargo.lock b/Cargo.lock index ace67d3a..4abd497a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -167,6 +167,15 @@ dependencies = [ "syn 2.0.116", ] +[[package]] +name = "atomic-polyfill" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cf2bce30dfe09ef0bfaef228b9d414faaf7e563035494d7fe092dba54b300f4" +dependencies = [ + "critical-section", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -250,15 +259,6 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" -[[package]] -name = "bincode" -version = "1.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" -dependencies = [ - "serde", -] - [[package]] name = "bitflags" version = "2.11.0" @@ -349,6 +349,12 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.11.1" @@ -502,6 +508,15 @@ dependencies = [ "cc", ] +[[package]] +name = "cobs" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fa961b519f0b462e3a3b4a34b64d119eeaca1d59af726fe450bbba07a9fc0a1" +dependencies = [ + "thiserror 2.0.18", +] + [[package]] name = "codespan-reporting" version = "0.13.1" @@ -614,6 +629,12 @@ dependencies = [ "itertools 0.10.5", ] +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -812,6 +833,18 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "embedded-io" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef1a6892d9eef45c8fa6b9e0086428a2cca8491aca8f787c534a3d6d0bcb3ced" + +[[package]] +name = "embedded-io" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d" + [[package]] name = "ember-client" version = "0.4.8" @@ -831,11 +864,11 @@ dependencies = [ name = "ember-cluster" version = "0.4.8" dependencies = [ - "bincode", "bytes", "crc32fast", "hmac", "openraft", + "postcard", "rand 0.9.2", "serde", "serde_json", @@ -932,6 +965,7 @@ dependencies = [ "clap", "colored", "dirs", + "ember-client", "ember-protocol", "rand 0.9.2", "rustls", @@ -1228,6 +1262,15 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hash32" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1258,6 +1301,20 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +[[package]] +name = "heapless" +version = "0.7.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdc6457c0eb62c71aac4bc17216026d8410337c4126773b9c5daba343f17964f" +dependencies = [ + "atomic-polyfill", + "hash32", + "rustc_version", + "serde", + "spin", + "stable_deref_trait", +] + [[package]] name = "heck" version = "0.5.0" @@ -1916,6 +1973,19 @@ version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "postcard" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6764c3b5dd454e283a30e6dfe78e9b31096d9e32036b5d1eaac7a6119ccb9a24" +dependencies = [ + "cobs", + "embedded-io 0.4.0", + "embedded-io 0.6.1", + "heapless", + "serde", +] + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -2347,6 +2417,15 @@ dependencies = [ "serde_json", ] +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "1.1.3" @@ -2651,6 +2730,21 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "static_assertions" version = "1.1.0" diff --git a/Cargo.toml b/Cargo.toml index 9b7e7c3d..0b968af5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,3 +95,4 @@ emberkv-core = { version = "0.4.8", path = "crates/ember-core" } ember-protocol = { version = "0.4.8", path = "crates/ember-protocol" } ember-persistence = { version = "0.4.8", path = "crates/ember-persistence" } ember-cluster = { version = "0.4.8", path = "crates/ember-cluster" } +ember-client = { version = "0.4.8", path = "crates/ember-client" } diff --git a/crates/ember-cli/Cargo.toml b/crates/ember-cli/Cargo.toml index 35e120e4..9e7ff293 100644 --- a/crates/ember-cli/Cargo.toml +++ b/crates/ember-cli/Cargo.toml @@ -18,6 +18,7 @@ clap = { workspace = true } tokio = { workspace = true } bytes = { workspace = true } ember-protocol = { workspace = true } +ember-client = { workspace = true } thiserror = { workspace = true } rand = { workspace = true } tokio-rustls = { workspace = true } diff --git a/crates/ember-cli/src/batch.rs b/crates/ember-cli/src/batch.rs index a536106e..e52015cf 100644 --- a/crates/ember-cli/src/batch.rs +++ b/crates/ember-cli/src/batch.rs @@ -13,8 +13,8 @@ use std::io::{self, BufRead}; use std::process::ExitCode; use colored::Colorize; +use ember_client::Client; -use crate::connection::Connection; use crate::format::format_response; use crate::tls::TlsClientConfig; @@ -46,7 +46,11 @@ async fn run_batch_async( password: Option<&str>, tls: Option<&TlsClientConfig>, ) -> ExitCode { - let mut conn = match Connection::connect(host, port, tls).await { + let conn = match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + }; + let mut conn = match conn { Ok(c) => c, Err(e) => { eprintln!( @@ -58,9 +62,9 @@ async fn run_batch_async( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("authentication failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } } @@ -71,7 +75,7 @@ async fn run_batch_async( Ok(l) => l, Err(e) => { eprintln!("{}", format!("read error: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } }; @@ -84,17 +88,18 @@ async fn run_batch_async( } let tokens: Vec = trimmed.split_whitespace().map(|s| s.to_string()).collect(); + let refs: Vec<&str> = tokens.iter().map(String::as_str).collect(); - match conn.send_command(&tokens).await { + match conn.send(&refs).await { Ok(frame) => println!("{}", format_response(&frame)), Err(e) => { eprintln!("{}", format!("error: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } } } - conn.shutdown().await; + conn.disconnect().await; ExitCode::SUCCESS } diff --git a/crates/ember-cli/src/bench_conn.rs b/crates/ember-cli/src/bench_conn.rs index fb0f2b6c..9616212c 100644 --- a/crates/ember-cli/src/bench_conn.rs +++ b/crates/ember-cli/src/bench_conn.rs @@ -11,7 +11,6 @@ use ember_protocol::parse::parse_frame; use ember_protocol::types::Frame; use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use crate::connection::auth_frame; use crate::tls::{self, MaybeTlsStream, TlsClientConfig}; /// Read buffer size for benchmark connections (256 KiB). @@ -50,7 +49,10 @@ impl BenchConnection { /// Authenticates with the server using AUTH. pub async fn authenticate(&mut self, password: &str) -> Result<(), String> { - let frame = auth_frame(password); + let frame = Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"AUTH")), + Frame::Bulk(Bytes::from(password.to_string())), + ]); let mut buf = BytesMut::new(); frame.serialize(&mut buf); diff --git a/crates/ember-cli/src/cluster.rs b/crates/ember-cli/src/cluster.rs index eddf46bd..c45919de 100644 --- a/crates/ember-cli/src/cluster.rs +++ b/crates/ember-cli/src/cluster.rs @@ -8,8 +8,8 @@ use std::process::ExitCode; use clap::Subcommand; use colored::Colorize; +use ember_client::Client; -use crate::connection::Connection; use crate::format::format_response; use crate::tls::TlsClientConfig; @@ -401,7 +401,10 @@ pub fn run_cluster( } rt.block_on(async { - let mut conn = match Connection::connect(host, port, tls).await { + let mut conn = match match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + } { Ok(c) => c, Err(e) => { eprintln!( @@ -413,15 +416,16 @@ pub fn run_cluster( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("authentication failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } } let tokens = cmd.to_tokens(); - let exit_code = match conn.send_command(&tokens).await { + let refs: Vec<&str> = tokens.iter().map(String::as_str).collect(); + let exit_code = match conn.send(&refs).await { Ok(frame) => { println!("{}", format_response(&frame)); ExitCode::SUCCESS @@ -432,7 +436,7 @@ pub fn run_cluster( } }; - conn.shutdown().await; + conn.disconnect().await; exit_code }) } @@ -462,7 +466,7 @@ async fn run_cluster_create( println!(">>> creating cluster with {primary_count} primaries and {replica_count} replicas"); // connect to all nodes and collect their IDs - let mut connections: Vec<(String, Connection, String)> = Vec::new(); // (addr, conn, node_id) + let mut connections: Vec<(String, Client, String)> = Vec::new(); // (addr, conn, node_id) for addr in addrs { let (host, port) = match parse_host_port(addr) { @@ -473,7 +477,10 @@ async fn run_cluster_create( } }; - let mut conn = match Connection::connect(&host, port, tls).await { + let mut conn = match match tls { + Some(tls) => Client::connect_tls(&host, port, tls).await, + None => Client::connect(&host, port).await, + } { Ok(c) => c, Err(e) => { eprintln!("{}", format!("could not connect to {addr}: {e}").red()); @@ -482,14 +489,14 @@ async fn run_cluster_create( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("auth failed on {addr}: {e}").red()); return ExitCode::FAILURE; } } // verify node is in cluster mode and has no slots assigned - match conn.send_command_strs(&["CLUSTER", "INFO"]).await { + match conn.send(&["CLUSTER", "INFO"]).await { Ok(frame) => { let info = frame_to_string(&frame); if !info.contains("cluster_state:") { @@ -520,7 +527,7 @@ async fn run_cluster_create( } // get node ID - let node_id = match conn.send_command_strs(&["CLUSTER", "MYID"]).await { + let node_id = match conn.send(&["CLUSTER", "MYID"]).await { Ok(frame) => frame_to_string(&frame), Err(e) => { eprintln!("{}", format!("CLUSTER MYID failed on {addr}: {e}").red()); @@ -547,7 +554,7 @@ async fn run_cluster_create( let start_s = start.to_string(); let end_s = end.to_string(); match conn - .send_command_strs(&["CLUSTER", "ADDSLOTSRANGE", &start_s, &end_s]) + .send(&["CLUSTER", "ADDSLOTSRANGE", &start_s, &end_s]) .await { Ok(frame) if is_ok(&frame) => { @@ -598,7 +605,7 @@ async fn run_cluster_create( match connections[0] .1 - .send_command_strs(&["CLUSTER", "MEET", &host, &port.to_string()]) + .send(&["CLUSTER", "MEET", &host, &port.to_string()]) .await { Ok(frame) if is_ok(&frame) => { @@ -625,7 +632,7 @@ async fn run_cluster_create( // also meet from other nodes back to the first node to speed convergence for (_, conn, _) in connections.iter_mut().skip(1) { let _ = conn - .send_command_strs(&["CLUSTER", "MEET", first_host, &first_port.to_string()]) + .send(&["CLUSTER", "MEET", first_host, &first_port.to_string()]) .await; } @@ -638,7 +645,7 @@ async fn run_cluster_create( let mut converged = true; for (addr, conn, _) in &mut connections { - match conn.send_command_strs(&["CLUSTER", "INFO"]).await { + match conn.send(&["CLUSTER", "INFO"]).await { Ok(frame) => { let info = frame_to_string(&frame); let known = parse_cluster_info_field(&info, "cluster_known_nodes") @@ -685,7 +692,7 @@ async fn run_cluster_create( match connections[replica_idx] .1 - .send_command_strs(&["CLUSTER", "REPLICATE", &primary_id]) + .send(&["CLUSTER", "REPLICATE", &primary_id]) .await { Ok(frame) if is_ok(&frame) => { @@ -720,7 +727,7 @@ async fn run_cluster_create( // clean up connections for (_, mut conn, _) in connections { - conn.shutdown().await; + conn.disconnect().await; } println!("{}", ">>> cluster created successfully".green()); @@ -748,7 +755,10 @@ async fn run_cluster_check( } }; - let mut conn = match Connection::connect(&host, port, tls).await { + let mut conn = match match tls { + Some(tls) => Client::connect_tls(&host, port, tls).await, + None => Client::connect(&host, port).await, + } { Ok(c) => c, Err(e) => { eprintln!("{}", format!("could not connect to {addr}: {e}").red()); @@ -757,33 +767,33 @@ async fn run_cluster_check( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("auth failed: {e}").red()); return ExitCode::FAILURE; } } // get cluster info - let cluster_info = match conn.send_command_strs(&["CLUSTER", "INFO"]).await { + let cluster_info = match conn.send(&["CLUSTER", "INFO"]).await { Ok(frame) => frame_to_string(&frame), Err(e) => { eprintln!("{}", format!("CLUSTER INFO failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } }; // get cluster nodes - let nodes_output = match conn.send_command_strs(&["CLUSTER", "NODES"]).await { + let nodes_output = match conn.send(&["CLUSTER", "NODES"]).await { Ok(frame) => frame_to_string(&frame), Err(e) => { eprintln!("{}", format!("CLUSTER NODES failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } }; - conn.shutdown().await; + conn.disconnect().await; // parse nodes let mut errors: Vec = Vec::new(); @@ -934,7 +944,7 @@ async fn run_cluster_reshard( }; // get cluster nodes to find addresses and slot ownership - let nodes_output = match conn.send_command_strs(&["CLUSTER", "NODES"]).await { + let nodes_output = match conn.send(&["CLUSTER", "NODES"]).await { Ok(frame) => frame_to_string(&frame), Err(e) => { eprintln!("{}", format!("CLUSTER NODES failed: {e}").red()); @@ -1031,7 +1041,7 @@ async fn run_cluster_reshard( for &slot in &slots_to_move { // 1. SETSLOT IMPORTING on target let r = target_conn - .send_command_strs(&[ + .send(&[ "CLUSTER", "SETSLOT", &slot.to_string(), @@ -1053,7 +1063,7 @@ async fn run_cluster_reshard( // 2. SETSLOT MIGRATING on source let r = source_conn - .send_command_strs(&[ + .send(&[ "CLUSTER", "SETSLOT", &slot.to_string(), @@ -1076,7 +1086,7 @@ async fn run_cluster_reshard( // 3. migrate keys loop { let keys_frame = match source_conn - .send_command_strs(&["CLUSTER", "GETKEYSINSLOT", &slot.to_string(), "100"]) + .send(&["CLUSTER", "GETKEYSINSLOT", &slot.to_string(), "100"]) .await { Ok(f) => f, @@ -1096,7 +1106,7 @@ async fn run_cluster_reshard( for key in &keys { let r = source_conn - .send_command_strs(&[ + .send(&[ "MIGRATE", &target_host, &target_port.to_string(), @@ -1124,7 +1134,7 @@ async fn run_cluster_reshard( let slot_s = slot.to_string(); for c in [&mut source_conn, &mut target_conn, &mut conn] { let _ = c - .send_command_strs(&["CLUSTER", "SETSLOT", &slot_s, "NODE", &target.id]) + .send(&["CLUSTER", "SETSLOT", &slot_s, "NODE", &target.id]) .await; } @@ -1134,9 +1144,9 @@ async fn run_cluster_reshard( } } - source_conn.shutdown().await; - target_conn.shutdown().await; - conn.shutdown().await; + source_conn.disconnect().await; + target_conn.disconnect().await; + conn.disconnect().await; if moved == slots_to_move.len() as u16 { println!( @@ -1183,14 +1193,14 @@ async fn run_cluster_rebalance( Err(code) => return code, }; - let nodes_output = match conn.send_command_strs(&["CLUSTER", "NODES"]).await { + let nodes_output = match conn.send(&["CLUSTER", "NODES"]).await { Ok(frame) => frame_to_string(&frame), Err(e) => { eprintln!("{}", format!("CLUSTER NODES failed: {e}").red()); return ExitCode::FAILURE; } }; - conn.shutdown().await; + conn.disconnect().await; let all_nodes = parse_nodes_output(&nodes_output); let primaries: Vec<&NodeInfo> = all_nodes.iter().filter(|n| n.is_primary).collect(); @@ -1302,14 +1312,18 @@ async fn run_cluster_rebalance( // helpers // --------------------------------------------------------------------------- -/// Connects to a node and authenticates. Returns the connection or an error exit code. +/// Connects to a node and authenticates. Returns the client or an error exit code. async fn connect_and_auth( host: &str, port: u16, password: Option<&str>, tls: Option<&TlsClientConfig>, -) -> Result { - let mut conn = Connection::connect(host, port, tls).await.map_err(|e| { +) -> Result { + let mut conn = match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + } + .map_err(|e| { eprintln!( "{}", format!("could not connect to {host}:{port}: {e}").red() @@ -1318,7 +1332,7 @@ async fn connect_and_auth( })?; if let Some(pw) = password { - conn.authenticate(pw).await.map_err(|e| { + conn.auth(pw).await.map_err(|e| { eprintln!("{}", format!("auth failed on {host}:{port}: {e}").red()); ExitCode::FAILURE })?; @@ -1402,9 +1416,7 @@ fn frame_to_string_list(frame: &ember_protocol::types::Frame) -> Vec { } /// Extracts a human-readable message from a connection result. -fn result_msg( - r: &Result, -) -> String { +fn result_msg(r: &Result) -> String { match r { Ok(f) => frame_to_string(f), Err(e) => e.to_string(), diff --git a/crates/ember-cli/src/connection.rs b/crates/ember-cli/src/connection.rs deleted file mode 100644 index 7b2d3320..00000000 --- a/crates/ember-cli/src/connection.rs +++ /dev/null @@ -1,193 +0,0 @@ -//! Async connection to an ember server (plain TCP or TLS). -//! -//! Handles connecting, sending commands as RESP3 arrays, -//! and reading back parsed frames. - -use std::time::Duration; - -use bytes::BytesMut; -use ember_protocol::parse::parse_frame; -use ember_protocol::types::Frame; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; - -use crate::tls::{self, MaybeTlsStream, TlsClientConfig}; - -/// Maximum read buffer size (64 KiB). Prevents unbounded memory growth -/// if the server sends a response that never completes. -const MAX_READ_BUF: usize = 64 * 1024; - -/// Default timeout for connecting to the server. -const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); - -/// Default timeout for reading a response from the server. -const READ_TIMEOUT: Duration = Duration::from_secs(10); - -/// Errors that can occur during connection operations. -#[derive(Debug, thiserror::Error)] -pub enum ConnectionError { - #[error("connection failed: {0}")] - Io(#[from] std::io::Error), - - #[error("protocol error: {0}")] - Protocol(String), - - #[error("server disconnected")] - Disconnected, - - #[error("authentication failed: {0}")] - AuthFailed(String), - - #[error("connection timed out")] - Timeout, - - #[error("response too large (exceeded {MAX_READ_BUF} bytes)")] - ResponseTooLarge, -} - -/// A connection to an ember server with read/write buffering. -/// -/// Works transparently over plain TCP or TLS — the underlying stream -/// type is determined at connect time. -pub struct Connection { - stream: MaybeTlsStream, - read_buf: BytesMut, - write_buf: BytesMut, -} - -impl Connection { - /// Connects to an ember server at the given host and port. - /// - /// When `tls` is `Some`, the connection is upgraded to TLS after the - /// TCP handshake. Times out after 5 seconds if the server is unreachable. - pub async fn connect( - host: &str, - port: u16, - tls: Option<&TlsClientConfig>, - ) -> Result { - let stream = tokio::time::timeout(CONNECT_TIMEOUT, tls::connect(host, port, tls)) - .await - .map_err(|_| ConnectionError::Timeout)? - .map_err(ConnectionError::Io)?; - - Ok(Self { - stream, - read_buf: BytesMut::with_capacity(4096), - write_buf: BytesMut::with_capacity(4096), - }) - } - - /// Sends a command (as a list of string tokens) and reads the response. - /// - /// Tokens are serialized as a RESP3 array of bulk strings, which is - /// the standard client -> server wire format. - pub async fn send_command(&mut self, tokens: &[String]) -> Result { - // build a RESP3 array of bulk strings - let parts: Vec = tokens - .iter() - .map(|t| Frame::Bulk(bytes::Bytes::from(t.clone()))) - .collect(); - let frame = Frame::Array(parts); - - // serialize into the write buffer and flush - self.write_buf.clear(); - frame.serialize(&mut self.write_buf); - self.stream.write_all(&self.write_buf).await?; - self.stream.flush().await?; - - // read the response - self.read_response().await - } - - /// Authenticates with the server using the AUTH command. - pub async fn authenticate(&mut self, password: &str) -> Result<(), ConnectionError> { - let frame = auth_frame(password); - self.write_buf.clear(); - frame.serialize(&mut self.write_buf); - self.stream.write_all(&self.write_buf).await?; - self.stream.flush().await?; - let response = self.read_response().await?; - - match &response { - Frame::Simple(s) if s == "OK" => Ok(()), - Frame::Error(e) => Err(ConnectionError::AuthFailed(e.clone())), - _ => Err(ConnectionError::AuthFailed( - "unexpected response to AUTH".into(), - )), - } - } - - /// Sends a command from a slice of string references and reads the response. - /// - /// Convenience wrapper around `send_command` for cases where tokens are - /// string slices rather than owned Strings. - pub async fn send_command_strs(&mut self, args: &[&str]) -> Result { - let owned: Vec = args.iter().map(|s| (*s).to_string()).collect(); - self.send_command(&owned).await - } - - /// Gracefully shuts down the connection. - /// - /// Sends a QUIT command to the server and then shuts down the stream. - /// Errors are intentionally ignored — this is best-effort cleanup. - pub async fn shutdown(&mut self) { - // try to send QUIT so the server can clean up - let quit = Frame::Array(vec![Frame::Bulk(bytes::Bytes::from_static(b"QUIT"))]); - self.write_buf.clear(); - quit.serialize(&mut self.write_buf); - let _ = self.stream.write_all(&self.write_buf).await; - let _ = self.stream.flush().await; - - // graceful shutdown (sends FIN instead of RST for TCP, or - // close_notify for TLS) - let _ = self.stream.shutdown().await; - } - - /// Reads a complete RESP3 frame from the server. - async fn read_response(&mut self) -> Result { - loop { - // try to parse a frame from what we have - if !self.read_buf.is_empty() { - match parse_frame(&self.read_buf) { - Ok(Some((frame, consumed))) => { - // remove the consumed bytes from the buffer - let _ = self.read_buf.split_to(consumed); - return Ok(frame); - } - Ok(None) => { - // incomplete — need more data - } - Err(e) => { - return Err(ConnectionError::Protocol(e.to_string())); - } - } - } - - // guard against unbounded buffer growth - if self.read_buf.len() >= MAX_READ_BUF { - return Err(ConnectionError::ResponseTooLarge); - } - - // read more data from the socket with a timeout - let read_result = - tokio::time::timeout(READ_TIMEOUT, self.stream.read_buf(&mut self.read_buf)).await; - - match read_result { - Ok(Ok(0)) => return Err(ConnectionError::Disconnected), - Ok(Ok(_)) => {} // got data, loop back to parse - Ok(Err(e)) => return Err(ConnectionError::Io(e)), - Err(_) => return Err(ConnectionError::Timeout), - } - } - } -} - -/// Builds an AUTH command frame from a password string. -/// -/// Shared by the interactive REPL connection and the benchmark connection so -/// the AUTH wire format is defined in exactly one place. -pub fn auth_frame(password: &str) -> Frame { - Frame::Array(vec![ - Frame::Bulk(bytes::Bytes::from_static(b"AUTH")), - Frame::Bulk(bytes::Bytes::from(password.to_string())), - ]) -} diff --git a/crates/ember-cli/src/main.rs b/crates/ember-cli/src/main.rs index 4000bc36..3d4b2de0 100644 --- a/crates/ember-cli/src/main.rs +++ b/crates/ember-cli/src/main.rs @@ -10,7 +10,6 @@ mod bench_conn; mod benchmark; mod cluster; mod commands; -mod connection; mod format; mod repl; mod tls; @@ -21,6 +20,7 @@ use std::process::ExitCode; use clap::{Parser, Subcommand}; use colored::Colorize; +use ember_client::Client; use crate::tls::TlsClientConfig; @@ -189,7 +189,10 @@ fn run_oneshot( }; rt.block_on(async { - let mut conn = match connection::Connection::connect(host, port, tls).await { + let mut conn = match match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + } { Ok(c) => c, Err(e) => { eprintln!( @@ -201,14 +204,15 @@ fn run_oneshot( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("authentication failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } } - let exit_code = match conn.send_command(command).await { + let refs: Vec<&str> = command.iter().map(String::as_str).collect(); + let exit_code = match conn.send(&refs).await { Ok(frame) => { println!("{}", format::format_response(&frame)); ExitCode::SUCCESS @@ -219,7 +223,7 @@ fn run_oneshot( } }; - conn.shutdown().await; + conn.disconnect().await; exit_code }) } diff --git a/crates/ember-cli/src/repl.rs b/crates/ember-cli/src/repl.rs index be1777b5..ad7d270b 100644 --- a/crates/ember-cli/src/repl.rs +++ b/crates/ember-cli/src/repl.rs @@ -16,10 +16,11 @@ use rustyline::hint::Hinter; use rustyline::validate::Validator; use rustyline::{CompletionType, Config, Context, Editor, Helper}; +use ember_client::{Client, ClientError}; + use crate::commands::{ command_names, commands_by_group, find_command, has_subcommands, subcommands, }; -use crate::connection::{Connection, ConnectionError}; use crate::format::format_response; use crate::tls::TlsClientConfig; @@ -45,7 +46,12 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC }; // connect to server - let mut conn = match rt.block_on(Connection::connect(host, port, tls)) { + let mut conn = match rt.block_on(async { + match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + } + }) { Ok(c) => c, Err(e) => { eprintln!( @@ -58,9 +64,9 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC // authenticate if needed if let Some(pw) = password { - if let Err(e) = rt.block_on(conn.authenticate(pw)) { + if let Err(e) = rt.block_on(conn.auth(pw)) { eprintln!("{}", format!("authentication failed: {e}").red()); - rt.block_on(conn.shutdown()); + rt.block_on(conn.disconnect()); return; } } @@ -74,7 +80,7 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC Ok(editor) => editor, Err(e) => { eprintln!("{}", format!("failed to create editor: {e}").red()); - rt.block_on(conn.shutdown()); + rt.block_on(conn.disconnect()); return; } }; @@ -133,11 +139,12 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC continue; } - match rt.block_on(conn.send_command(&tokens)) { + let refs: Vec<&str> = tokens.iter().map(String::as_str).collect(); + match rt.block_on(conn.send(&refs)) { Ok(frame) => { println!("{}", format_response(&frame)); } - Err(ConnectionError::Disconnected) => { + Err(ClientError::Disconnected) => { eprintln!("{}", "server disconnected, reconnecting...".yellow()); match rt.block_on(reconnect(host, port, password, tls)) { Ok(new_conn) => { @@ -182,7 +189,7 @@ pub fn run_repl(host: &str, port: u16, password: Option<&str>, tls: Option<&TlsC } // graceful shutdown — send QUIT and close the TCP stream - rt.block_on(conn.shutdown()); + rt.block_on(conn.disconnect()); } /// Establishes a new connection, authenticating if a password is provided. @@ -191,10 +198,13 @@ async fn reconnect( port: u16, password: Option<&str>, tls: Option<&TlsClientConfig>, -) -> Result { - let mut conn = Connection::connect(host, port, tls).await?; +) -> Result { + let mut conn = match tls { + Some(tls) => Client::connect_tls(host, port, tls).await?, + None => Client::connect(host, port).await?, + }; if let Some(pw) = password { - conn.authenticate(pw).await?; + conn.auth(pw).await?; } Ok(conn) } diff --git a/crates/ember-cli/src/tls.rs b/crates/ember-cli/src/tls.rs index decd036a..1dea80fc 100644 --- a/crates/ember-cli/src/tls.rs +++ b/crates/ember-cli/src/tls.rs @@ -1,8 +1,8 @@ //! TLS client support for ember-cli. //! //! Provides a `MaybeTlsStream` wrapper that implements `AsyncRead` and -//! `AsyncWrite`, allowing the rest of the codebase to work with either -//! plain TCP or TLS connections transparently. +//! `AsyncWrite`, allowing `bench_conn` to work directly with either plain +//! TCP or TLS streams. Interactive commands use `ember_client::Client` instead. use std::io; use std::pin::Pin; @@ -16,17 +16,9 @@ use tokio::net::TcpStream; use tokio_rustls::client::TlsStream; use tokio_rustls::TlsConnector; -/// TLS configuration for client connections. -#[derive(Clone, Debug)] -pub struct TlsClientConfig { - /// Optional path to a CA certificate (PEM) for verifying the server. - /// When `None`, the system trust store is used. - pub ca_cert: Option, - - /// Skip server certificate verification entirely. Prints a warning - /// to stderr when enabled — useful for development with self-signed certs. - pub insecure: bool, -} +/// TLS client configuration — re-exported from `ember-client` so all of the +/// CLI uses one canonical type. +pub use ember_client::TlsClientConfig; /// A TCP stream that may or may not be wrapped in TLS. pub enum MaybeTlsStream { diff --git a/crates/ember-cli/src/watch.rs b/crates/ember-cli/src/watch.rs index f24311a5..62ce1101 100644 --- a/crates/ember-cli/src/watch.rs +++ b/crates/ember-cli/src/watch.rs @@ -9,8 +9,8 @@ use std::process::ExitCode; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use colored::Colorize; +use ember_client::Client; -use crate::connection::Connection; use crate::format::format_response; use crate::tls::TlsClientConfig; @@ -43,7 +43,11 @@ async fn run_watch_async( key: &str, interval_ms: u64, ) -> ExitCode { - let mut conn = match Connection::connect(host, port, tls).await { + let conn = match tls { + Some(tls) => Client::connect_tls(host, port, tls).await, + None => Client::connect(host, port).await, + }; + let mut conn = match conn { Ok(c) => c, Err(e) => { eprintln!( @@ -55,9 +59,9 @@ async fn run_watch_async( }; if let Some(pw) = password { - if let Err(e) = conn.authenticate(pw).await { + if let Err(e) = conn.auth(pw).await { eprintln!("{}", format!("authentication failed: {e}").red()); - conn.shutdown().await; + conn.disconnect().await; return ExitCode::FAILURE; } } @@ -84,7 +88,7 @@ async fn run_watch_async( } _ = tokio::time::sleep(interval) => { - match conn.send_command_strs(&["GET", key]).await { + match conn.send(&["GET", key]).await { Ok(frame) => { let formatted = format_response(&frame); if Some(&formatted) != last_value.as_ref() { @@ -102,7 +106,7 @@ async fn run_watch_async( } } - conn.shutdown().await; + conn.disconnect().await; exit_code } diff --git a/crates/ember-client/src/lib.rs b/crates/ember-client/src/lib.rs index 2081eedc..2b9d32ed 100644 --- a/crates/ember-client/src/lib.rs +++ b/crates/ember-client/src/lib.rs @@ -23,3 +23,5 @@ pub mod tls; pub use connection::{Client, ClientError}; pub use ember_protocol::types::Frame; +#[cfg(feature = "tls")] +pub use tls::TlsClientConfig;