diff --git a/Cargo.toml b/Cargo.toml index 4bdd8c8..8898805 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,4 +6,4 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] -#tokio = { version = "1", features = ["full"] } +byteorder = "1.4" # For byte manipulation in a cross-platform manner diff --git a/src/main.rs b/src/main.rs index 055bca8..4fb2040 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,13 +5,13 @@ use std::net::TcpListener; use std::path::Path; use std::thread; use std::process; +use byteorder::{ByteOrder, BigEndian}; fn main() { let args: Vec = env::args().collect(); let mut port = "7890"; // default port let mut log_dir = "/var/log/RMS-server"; // default log directory - // Command line argument parsing if args.len() > 1 { let mut i = 1; while i < args.len() { @@ -19,22 +19,21 @@ fn main() { "-p" | "--port" => { if i + 1 < args.len() { port = &args[i + 1]; - i += 1; + i += 2; } - } + }, "-l" => { if i + 1 < args.len() { log_dir = &args[i + 1]; - i += 1; + i += 2; } - } + }, "-h" | "--help" => { print_usage(); return; - } - _ => {} + }, + _ => i += 1, } - i += 1; } } @@ -43,50 +42,65 @@ fn main() { process::exit(1); }); - // Ensure the log directory exists std::fs::create_dir_all(log_dir).expect("Failed to create log directory"); - log_verbose(format!("Server listening on port {}", port), log_dir); for stream in listener.incoming() { match stream { Ok(stream) => { - let peer_addr = stream.peer_addr().unwrap(); // It's safe to unwrap here as connection is already established + let peer_addr = stream.peer_addr().unwrap(); log_verbose(format!("Client connected: {}", peer_addr), log_dir); - let log_dir = log_dir.to_string(); // Clone log_dir for the thread + let log_dir = log_dir.to_string(); thread::spawn(move || { handle_client(stream, &log_dir); }); - } - Err(e) => { - log_verbose(format!("Connection failed: {}", e), log_dir); - } + }, + Err(e) => log_verbose(format!("Connection failed: {}", e), log_dir), } } } -// Function to handle a single client fn handle_client(mut stream: std::net::TcpStream, log_dir: &str) { - let mut buffer = [0; 1024]; + println!("Handling new client"); + let mut buffer = Vec::new(); + let mut data = vec![0u8; 1024]; loop { - match stream.read(&mut buffer) { + match stream.read(&mut data) { Ok(0) => { - break; // Client disconnected + println!("Client disconnected"); + break; }, Ok(size) => { - let msg = String::from_utf8_lossy(&buffer[..size]); - println!("{}", msg); // Print what client sends - log_message(msg.into_owned(), log_dir); // Log to messages file + println!("Received {} bytes", size); + buffer.extend_from_slice(&data[..size]); + + if buffer.len() < 5 { + continue; // Wait for enough bytes to read header + } + + let msg_type = buffer[0]; + let msg_len = BigEndian::read_u32(&buffer[1..5]) as usize; + + if buffer.len() < 5 + msg_len { + continue; // Wait for the full message + } + + let msg = String::from_utf8_lossy(&buffer[5..5+msg_len]); + println!("Type: {}, Msg: {}", msg_type, msg); + log_message(msg.to_string(), log_dir); + // Clear the buffer or handle the next message + buffer.drain(..5+msg_len); }, - Err(_) => { - break; // Connection error + Err(e) => { + println!("Failed to read from client: {:?}", e); + break; } } } } -// Log verbose information to a file + fn log_verbose(message: String, log_dir: &str) { let mut file = OpenOptions::new() .create(true) @@ -94,27 +108,22 @@ fn log_verbose(message: String, log_dir: &str) { .append(true) .open(Path::new(log_dir).join("verbose.log")) .expect("Failed to open verbose log file"); - writeln!(file, "{}", message).expect("Failed to write to verbose log file"); } -// Log raw messages to a separate file fn log_message(message: String, log_dir: &str) { let mut file = OpenOptions::new() .create(true) .write(true) .append(true) - .open(Path::new(log_dir).join("rawmessages.log")) + .open(Path::new(log_dir).join("messages.log")) .expect("Failed to open message log file"); - writeln!(file, "{}", message).expect("Failed to write to message log file"); } -// Print usage information fn print_usage() { println!("Usage: rms-server [-p/--port PORT] [-l LOG_DIR] [-h/--help]"); println!("Options:"); println!(" -p, --port PORT Set the port number (default: 7890)"); - println!(" -l LOG_DIR Set the log directory (default: /var/log/RMS-server)"); - println!(" -h, --help Print this help message"); + println!(" -l LOG_DIR"); }