148 lines
4.2 KiB
Rust
148 lines
4.2 KiB
Rust
use std::io::{self, Write};
|
|
use std::net::TcpStream;
|
|
use std::time::Duration;
|
|
use std::env;
|
|
use colored::Colorize;
|
|
use regex::Regex;
|
|
use std::process;
|
|
|
|
mod resp;
|
|
|
|
const DEFAULT_PORT: u16 = 9880;
|
|
const PROMPT_NAME: &str = "futriix";
|
|
const CONNECTION_TIMEOUT_SECS: u64 = 2;
|
|
|
|
fn main() {
|
|
let (host, port) = parse_args();
|
|
let addr = format!("{}:{}", host, port);
|
|
|
|
let stream = match TcpStream::connect_timeout(
|
|
&addr.parse().unwrap(),
|
|
Duration::from_secs(CONNECTION_TIMEOUT_SECS)
|
|
) {
|
|
Ok(stream) => stream,
|
|
Err(e) if e.kind() == io::ErrorKind::ConnectionRefused => {
|
|
eprintln!("{}", "Connection refused".red());
|
|
process::exit(1);
|
|
},
|
|
Err(e) => {
|
|
eprintln!("Connection error: {}", e);
|
|
process::exit(1);
|
|
}
|
|
};
|
|
|
|
println!("Connected to {}", addr.green());
|
|
run_repl_loop(stream, &host, port);
|
|
}
|
|
|
|
fn parse_args() -> (String, u16) {
|
|
let args: Vec<String> = env::args().collect();
|
|
let mut host = "127.0.0.1".to_string();
|
|
let mut port = DEFAULT_PORT;
|
|
|
|
if args.len() > 1 {
|
|
let re = Regex::new(r"^(?:([^:]+):)?([^:]+)(?::(\d+))?$").unwrap();
|
|
if let Some(caps) = re.captures(&args[1]) {
|
|
if let Some(h) = caps.get(2) {
|
|
host = h.as_str().to_string();
|
|
}
|
|
if let Some(p) = caps.get(3) {
|
|
port = p.as_str().parse().unwrap_or(DEFAULT_PORT);
|
|
}
|
|
}
|
|
}
|
|
(host, port)
|
|
}
|
|
|
|
fn run_repl_loop(stream: TcpStream, host: &str, port: u16) {
|
|
let mut input = String::new();
|
|
loop {
|
|
print_prompt(host, port);
|
|
input.clear();
|
|
|
|
if io::stdin().read_line(&mut input).is_err() {
|
|
eprintln!("{}", "Failed to read input".red());
|
|
continue;
|
|
}
|
|
|
|
let input = input.trim();
|
|
if input.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
if input.eq_ignore_ascii_case("quit") || input.eq_ignore_ascii_case("exit") {
|
|
break;
|
|
}
|
|
|
|
if !is_valid_command(input) {
|
|
eprintln!("{}", "Error: Invalid command format".red());
|
|
continue;
|
|
}
|
|
|
|
match send_command(&stream, input) {
|
|
Ok(response) => print_response(&response),
|
|
Err(e) => {
|
|
if is_connection_error(&e) {
|
|
eprintln!("{}", "Connection error".red());
|
|
break;
|
|
}
|
|
eprintln!("{}", format!("Error: {}", e.to_string().replace("KeyDB", "Futriix")).red());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn print_prompt(host: &str, port: u16) {
|
|
let prompt = format!("{}:{}:{}:~>", PROMPT_NAME, host, port);
|
|
print!("{} ", prompt.green());
|
|
io::stdout().flush().unwrap();
|
|
}
|
|
|
|
fn is_valid_command(cmd: &str) -> bool {
|
|
!cmd.is_empty() &&
|
|
!cmd.chars().any(|c| c.is_control()) &&
|
|
cmd.split_whitespace().next().is_some()
|
|
}
|
|
|
|
fn send_command(stream: &TcpStream, command: &str) -> io::Result<resp::Value> {
|
|
let parts: Vec<&str> = command.split_whitespace().collect();
|
|
let mut resp_command = format!("*{}\r\n", parts.len());
|
|
|
|
for part in parts {
|
|
resp_command.push_str(&format!("${}\r\n{}\r\n", part.len(), part));
|
|
}
|
|
|
|
let mut stream = stream.try_clone()?;
|
|
stream.write_all(resp_command.as_bytes())?;
|
|
|
|
let mut decoder = resp::Decoder::new(&stream);
|
|
decoder.decode()
|
|
}
|
|
|
|
fn print_response(value: &resp::Value) {
|
|
match value {
|
|
resp::Value::SimpleString(s) | resp::Value::BulkString(s) => {
|
|
println!("{}", s.replace("KeyDB", "Futriix"));
|
|
},
|
|
resp::Value::Error(e) => {
|
|
println!("{}", format!("(error) {}", e.replace("KeyDB", "Futriix")).red());
|
|
},
|
|
resp::Value::Integer(i) => println!("(integer) {}", i),
|
|
resp::Value::Array(arr) => {
|
|
for (i, item) in arr.iter().enumerate() {
|
|
print!("{}) ", i + 1);
|
|
print_response(item);
|
|
}
|
|
},
|
|
resp::Value::Null => println!("(nil)"),
|
|
}
|
|
}
|
|
|
|
fn is_connection_error(error: &io::Error) -> bool {
|
|
matches!(
|
|
error.kind(),
|
|
io::ErrorKind::ConnectionAborted |
|
|
io::ErrorKind::ConnectionReset |
|
|
io::ErrorKind::BrokenPipe
|
|
)
|
|
} |