diff --git a/src/console.rs b/src/console.rs new file mode 100644 index 0000000..c1d0f42 --- /dev/null +++ b/src/console.rs @@ -0,0 +1,45 @@ +use std::sync::mpsc::Receiver; + +use crate::rcon::MAX_CONTENT_SIZE; + +pub struct ConsoleAccess { + console_recv: Receiver, + cmd_buffer: Vec, +} + +impl ConsoleAccess { + pub fn new(recv: Receiver) -> Self { + Self { + console_recv: recv, + cmd_buffer: Vec::new(), + } + } + + pub fn next_line(&self) -> Option { + self.console_recv.try_recv().ok() + } + + pub fn next_line_catpure(&mut self) -> Option { + if let Some(line) = self.next_line() { + let line_size = line.len(); + let mut buffer_size = 0; + + for bline in self.cmd_buffer.drain(..).rev().collect::>() { + buffer_size += bline.len(); + if buffer_size + line_size > MAX_CONTENT_SIZE { + break; + } + + self.cmd_buffer.insert(0, bline); + } + self.cmd_buffer.push(line.clone()); + + return Some(line); + } + None + } + + pub fn get_last_console_output(&self) -> &[String] { + &self.cmd_buffer + } +} diff --git a/src/console_hook.rs b/src/console_hook.rs index 3ad04da..310e854 100644 --- a/src/console_hook.rs +++ b/src/console_hook.rs @@ -22,7 +22,7 @@ static_detour! { pub fn hook_write_console() { unsafe { if !std::env::args().any(|arg| arg == "-dedicated") { - log::info!("this isn't a dedicated server; not hooking WriteConsoleA"); + log::warn!("this isn't a dedicated server; not hooking WriteConsoleA"); return; } @@ -42,7 +42,7 @@ pub fn hook_write_console() { pub fn hook_console_print(addr: isize) -> Option<()> { unsafe { - if PLUGIN.wait().console_recv.try_lock().is_some() { + if PLUGIN.wait().server.is_none() { log::warn!("rcon not running -> no Print hook"); return None; } diff --git a/src/lib.rs b/src/lib.rs index f414d4e..4d48846 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,37 +1,31 @@ -use bindings::{CmdSource, EngineFunctions, ENGINE_FUNCTIONS}; +use bindings::{EngineFunctions, ENGINE_FUNCTIONS}; use console_hook::{hook_console_print, hook_write_console}; use parking_lot::Mutex; -use rcon::{RconServer, RconTask}; -use rrplug::{ - bindings::class_types::client::CClient, mid::engine::WhichDll, prelude::*, to_c_string, -}; +use rcon::RconServer; +use rrplug::{mid::engine::WhichDll, prelude::*}; use std::{ collections::HashMap, env, - ffi::c_void, - sync::mpsc::{self, Receiver, Sender}, + sync::mpsc::{self, Sender}, }; pub mod bindings; +pub mod console; pub mod console_hook; pub mod rcon; const VALID_RCON_ARGS: [&str; 2] = ["rcon_ip_port", "rcon_password"]; -#[derive(Debug)] pub struct RconPlugin { - rcon_tasks: Mutex>, - rcon_send_tasks: Mutex>, // mutex is not needed but it must sync so clone on each thread console_sender: Mutex>, - console_recv: Mutex>, + server: Option>, } impl Plugin for RconPlugin { fn new(_: &PluginData) -> Self { - let (sender, recv) = mpsc::channel(); let (console_sender, console_recv) = mpsc::channel(); - let args = env::args() + let rcon_args = env::args() .zip(env::args().skip(1)) .filter(|(cmd, _)| VALID_RCON_ARGS.contains(&&cmd[..])) .fold(HashMap::new(), |mut hash_map, (cmd, arg)| { @@ -39,80 +33,43 @@ impl Plugin for RconPlugin { hash_map }); - std::thread::spawn(move || _ = run_rcon(args)); + let mut server = None; + + 'start_server: { + let (Some(bind_ip), Some(password)) = ( + rcon_args.get(VALID_RCON_ARGS[0]), + rcon_args.get(VALID_RCON_ARGS[1]), + ) else { + log::error!("the rcon args that were provided are invalid!"); + break 'start_server; + }; + + server = RconServer::try_new(&bind_ip, password, console_recv) + .map_err(|err| log::info!("failed to connect to socket : {err:?}")) + .map(|s| { + hook_write_console(); + s + }) + .ok(); + } Self { - rcon_tasks: Mutex::new(recv), - rcon_send_tasks: Mutex::new(sender), console_sender: Mutex::new(console_sender), - console_recv: Mutex::new(console_recv), + server: server.map(|s| s.into()), } } - fn on_dll_load(&self, engine: Option<&EngineData>, dll_ptr: &DLLPointer) { + fn on_dll_load(&self, _: Option<&EngineData>, dll_ptr: &DLLPointer) { unsafe { EngineFunctions::try_init(dll_ptr, &ENGINE_FUNCTIONS) }; if let WhichDll::Client = dll_ptr.which_dll() { let addr = dll_ptr.get_dll_ptr() as isize; std::thread::spawn(move || _ = hook_console_print(addr)); } - - let engine = if let Some(engine) = engine { - engine - } else { - return; - }; } fn runframe(&self) { - // can be moved somewhere else - - let funcs = ENGINE_FUNCTIONS.wait(); - - if let Ok(task) = self.rcon_tasks.lock().try_recv() { - match task { - RconTask::Runcommand(cmd) => unsafe { - log::info!("executing command : {cmd}"); - - let cmd = to_c_string!(cmd); - (funcs.cbuf_add_text_type)( - (funcs.cbuf_get_current_player)(), - cmd.as_ptr(), - CmdSource::Code, - ); - }, - } - } - } -} - -fn run_rcon(rcon_args: HashMap) -> Option { - let mut server = match RconServer::try_new( - rcon_args.get(VALID_RCON_ARGS[0])?, - rcon_args.get(VALID_RCON_ARGS[1])?.to_string(), - ) { - Ok(sv) => sv, - Err(err) => { - log::info!("failed to connect to socket : {err:?}"); - return None; - } - }; - - hook_write_console(); - - let rcon = PLUGIN.wait(); - - let rcon_send_tasks = rcon.rcon_send_tasks.lock(); - let console_recv = rcon.console_recv.lock(); - - loop { - let new_console_line = console_recv.try_recv().ok(); - - if let Some(tasks) = server.run(new_console_line) { - tasks - .into_iter() - .for_each(|task| rcon_send_tasks.send(task).expect("failed to send tasks")) - } + _ = self.server.as_ref().map(|s| s.lock().run()); } } diff --git a/src/rcon.rs b/src/rcon.rs index 644f7af..1ee110a 100644 --- a/src/rcon.rs +++ b/src/rcon.rs @@ -1,15 +1,15 @@ +use rrplug::to_c_string; use std::{ io::{self, Read, Write}, net::{TcpListener, TcpStream}, + sync::mpsc::Receiver, }; use thiserror::Error; -#[derive(Debug, Clone)] -pub enum RconTask { - Runcommand(String), -} - -// this could be fun https://discord.com/channels/920776187884732556/922663696273125387/1134900622773194782 +use crate::{ + bindings::{CmdSource, ENGINE_FUNCTIONS}, + console::ConsoleAccess, +}; const SERVERDATA_AUTH: i32 = 3; const SERVERDATA_EXECCOMMAND: i32 = 2; @@ -17,7 +17,7 @@ const SERVERDATA_AUTH_RESPONSE: i32 = 2; const SERVERDATA_RESPONSE_VALUE: i32 = 0; const MAX_PACKET_SIZE: usize = 4096; const MIN_PACKET_SIZE: usize = 10; -const MAX_CONTENT_SIZE: usize = MAX_PACKET_SIZE - MIN_PACKET_SIZE; +pub const MAX_CONTENT_SIZE: usize = MAX_PACKET_SIZE - MIN_PACKET_SIZE; #[derive(Debug, Error)] pub enum RconRequestError { @@ -80,42 +80,31 @@ pub struct RconServer { password: String, server: TcpListener, connections: Vec, - tasks: Vec, - cmd_buffer: Vec, + console: ConsoleAccess, } impl RconServer { - pub fn try_new(bind_ip: &str, password: String) -> Result { + pub fn try_new( + bind_ip: &str, + password: impl Into, + console_recv: Receiver, + ) -> Result { let server = TcpListener::bind(bind_ip)?; server.set_nonblocking(true)?; let rcon_server = Self { - password, + password: password.into(), server, connections: Vec::new(), - tasks: Vec::new(), - cmd_buffer: Vec::new(), + console: ConsoleAccess::new(console_recv), }; Ok(rcon_server) } - pub fn run(&mut self, new_console_line: Option) -> Option> { - if let Some(line) = new_console_line { - let line_size = line.len(); - let mut buffer_size = 0; - - for bline in self.cmd_buffer.drain(..).rev().collect::>() { - buffer_size += bline.len(); - if buffer_size + line_size > MAX_CONTENT_SIZE { - break; - } - - self.cmd_buffer.insert(0, bline); - } - self.cmd_buffer.push(line); - } + pub fn run(&mut self) { + while let Some(_) = self.console.next_line_catpure() {} // string allocation could be remove match self.server.accept() { Ok((conn, addr)) => match conn.set_nonblocking(true) { @@ -135,12 +124,8 @@ impl RconServer { } for i in 0..self.connections.len() { - match handle_connection(&mut self.connections[i], &self.password, &self.cmd_buffer) { - Ok(maybe_task) => { - if let Some(task) = maybe_task { - self.tasks.push(task) - } - } + match handle_connection(&mut self.connections[i], &self.password, &mut self.console) { + Ok(_) => {} Err(err) => { match &err { RconRequestError::SocketError(err) => match err.kind() { @@ -156,22 +141,25 @@ impl RconServer { } } } - - if self.tasks.is_empty() { - None - } else { - Some(self.tasks.drain(0..).collect()) - } } } pub fn handle_connection( conn: &mut RconStream, password: &str, - cmd_buffer: &[String], -) -> Result, RconRequestError> { - let stream = &mut conn.stream; + console: &mut ConsoleAccess, +) -> Result<(), RconRequestError> { + let (client_id, request_type, content) = read_rcon_stream(&mut conn.stream)?; + + let response = parse_response(conn, password, console, client_id, request_type, content)?; + + let buf: Vec = response.into(); + conn.stream.write_all(&buf)?; + Ok(()) +} + +fn read_rcon_stream(stream: &mut TcpStream) -> Result<(i32, i32, String), RconRequestError> { let mut size_buf = vec![0; 4]; let bytes_read = stream.read(&mut size_buf)?; if bytes_read != 4 { @@ -209,56 +197,87 @@ pub fn handle_connection( .or(Err(RconRequestError::VecToArrayError))?, ); - let content = String::from_utf8_lossy(&buf).to_string().replace('\0', ""); + Ok(( + client_id, + request_type, + String::from_utf8_lossy(&buf).to_string().replace('\0', ""), + )) // TODO: maybe use a struct +} - let (response, task) = match request_type { +fn parse_response( + conn: &mut RconStream, + password: &str, + console: &mut ConsoleAccess, + client_id: i32, + request_type: i32, + content: String, +) -> Result { + let response = match request_type { SERVERDATA_AUTH => { if content == password { conn.auth = true; log::info!("auth successful"); - ( - RconResponse { - id: client_id, - ty: SERVERDATA_AUTH_RESPONSE, - content: String::new(), - }, - None, - ) + RconResponse { + id: client_id, + ty: SERVERDATA_AUTH_RESPONSE, + content: String::new(), + } } else { log::warn!("auth failed"); conn.auth = false; - ( - RconResponse { - id: -1, - ty: SERVERDATA_AUTH_RESPONSE, - content: String::new(), - }, - None, - ) + RconResponse { + id: -1, + ty: SERVERDATA_AUTH_RESPONSE, + content: String::new(), + } + } + } + SERVERDATA_EXECCOMMAND if content == "dumpconsole" => { + if !conn.auth { + Err(RconRequestError::InvalidClientID(client_id))? + } + log::info!("sending console dump"); + + RconResponse { + id: client_id, + ty: SERVERDATA_RESPONSE_VALUE, + content: console.get_last_console_output().iter().cloned().collect(), } } SERVERDATA_EXECCOMMAND => { if !conn.auth { Err(RconRequestError::InvalidClientID(client_id))? } + log::info!("executing command : {content}"); + + let cmd = to_c_string!(content); + let funcs = ENGINE_FUNCTIONS.wait(); + unsafe { + (funcs.cbuf_add_text_type)( + (funcs.cbuf_get_current_player)(), + cmd.as_ptr(), + CmdSource::Code, + ); + + (funcs.cbuf_execute)() // execute the buffer rn since we want the results immediately + } - ( - RconResponse { - id: client_id, - ty: SERVERDATA_RESPONSE_VALUE, - content: cmd_buffer.iter().cloned().collect(), - }, - Some(RconTask::Runcommand(content)), - ) + let mut response = String::new(); + while let Some(console_out) = console.next_line_catpure() { + response += &console_out; + } + + RconResponse { + id: client_id, + ty: SERVERDATA_RESPONSE_VALUE, + content: response, + } } request_num => Err(RconRequestError::InvalidRequestType(request_num))?, }; - let buf: Vec = response.into(); - stream.write_all(&buf)?; - - Ok(task) + Ok(response) }