From 6abbe4d28210004b128955ff715ca1684d654af2 Mon Sep 17 00:00:00 2001 From: HoLLy Date: Fri, 13 Nov 2020 23:39:14 +0100 Subject: [PATCH] Allow creating handlers to have different behavior --- src/main.rs | 3 +- src/server/commands_request.rs | 20 +---- src/server/handler.rs | 40 ++++++++++ src/server/mod.rs | 4 +- src/server/server.rs | 134 ++++++++++++++++++--------------- 5 files changed, 118 insertions(+), 83 deletions(-) create mode 100644 src/server/handler.rs diff --git a/src/main.rs b/src/main.rs index 6ea12c7..c7d348e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,8 +8,7 @@ fn main() { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { - let server = server::CheatEngineServer::new(); - server.run().await.unwrap(); + server::run::().await.unwrap(); }) } diff --git a/src/server/commands_request.rs b/src/server/commands_request.rs index c385d91..f19d8b2 100644 --- a/src/server/commands_request.rs +++ b/src/server/commands_request.rs @@ -1,17 +1,15 @@ use std::fmt::Debug; use bytes::Buf; -use log::warn; use super::{command::*, commands_response::*}; use super::command::Command; -pub trait CERequest : Debug { - type Response: CEResponse + Debug; +pub trait CERequest : Debug + Send { + type Response: CEResponse + Debug + Send; const ID: Command; fn read(buf: &mut dyn Buf) -> Self; - fn process(self) -> Self::Response; } #[derive(Debug)] @@ -31,13 +29,6 @@ impl CERequest for CreateToolHelp32SnapshotRequest { process_id: buf.get_u32_le(), } } - - fn process(self) -> CreateToolHelp32SnapshotResponse { - warn!("Stubbed CreateToolHelp32SnapshotResponse::process"); // TODO - CreateToolHelp32SnapshotResponse { - handle: 0 - } - } } #[derive(Debug)] @@ -55,11 +46,4 @@ impl CERequest for Process32FirstRequest { handle: buf.get_u32_le(), } } - - fn process(self) -> Process32FirstResponse { - warn!("Stubbed Process32FirstResponse::process"); // TODO - Process32FirstResponse { - result: false, - } - } } diff --git a/src/server/handler.rs b/src/server/handler.rs new file mode 100644 index 0000000..703c61e --- /dev/null +++ b/src/server/handler.rs @@ -0,0 +1,40 @@ +use log::warn; + +use super::{commands_response::*, commands_request::*}; + +pub trait HandlerFactory { +} + +pub trait FullHandler + : Handler + + Handler { + fn create() -> Self; +} + +pub trait Handler { + fn handle(&self, req: TReq) -> TReq::Response; +} + +pub struct TestHandler; + +impl FullHandler for TestHandler { + fn create() -> TestHandler { TestHandler } +} + +impl Handler for TestHandler { + fn handle(&self, _req: CreateToolHelp32SnapshotRequest) -> CreateToolHelp32SnapshotResponse { + warn!("Stubbed CreateToolHelp32SnapshotResponse::process"); // TODO + CreateToolHelp32SnapshotResponse { + handle: 0 + } + } +} + +impl Handler for TestHandler { + fn handle(&self, _req: Process32FirstRequest) -> Process32FirstResponse { + warn!("Stubbed Process32FirstResponse::process"); // TODO + Process32FirstResponse { + result: false, + } + } +} diff --git a/src/server/mod.rs b/src/server/mod.rs index e177374..ab42403 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -2,5 +2,7 @@ mod server; mod command; mod commands_request; mod commands_response; +mod handler; -pub use server::CheatEngineServer; \ No newline at end of file +pub use server::run; +pub use handler::TestHandler; \ No newline at end of file diff --git a/src/server/server.rs b/src/server/server.rs index 246042a..44bc15d 100644 --- a/src/server/server.rs +++ b/src/server/server.rs @@ -1,79 +1,89 @@ -use log::{info, debug, error}; -use tokio::net::TcpListener; -use tokio::prelude::*; use bytes::Buf; +use log::{info, debug, error}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::prelude::*; -use super::{commands_request::*, commands_response::*}; +use super::{handler::*, commands_request::*, commands_response::*}; const DEFAULT_PORT: u16 = 52736; -pub struct CheatEngineServer { - port: u16 +pub async fn run() -> Result<(), Box> { + let listener = TcpListener::bind(("127.0.0.1", DEFAULT_PORT)).await?; + + loop { + let (socket, addr) = listener.accept().await?; + info!("New socket opened: {}", addr); + + tokio::spawn(async move { + let conn = CheatEngineConnection(socket); + let handler = T::create(); + conn.run(handler).await; + }); + } } -impl CheatEngineServer { - pub fn new() -> Self { - CheatEngineServer::new_port(DEFAULT_PORT) - } +pub struct CheatEngineConnection(TcpStream); - pub fn new_port(port: u16) -> Self { - CheatEngineServer {port} - } +// generics werent powerful enough +macro_rules! gen_request_dispatch { + ($bytes: ident for $handler: ident, $($request: ty),*) => { + match $bytes.get_u8() { + $( + <$request>::ID => gen_request_dispatch!($request, $bytes, $handler), + )* + byte => todo!("Unimplemented packet byte {}", byte), + } + }; + ($request: ty, $bytes: ident, $handler: ident) => { + { + let mut buffer = Vec::new(); - pub async fn run(&self) -> Result<(), Box> { - let listener = TcpListener::bind(("127.0.0.1", self.port)).await?; + let request = <$request>::read(&mut $bytes); + + debug!("Received item {:?}", request); + let response = $handler.handle(request); + debug!("... responding with {:?}", response); + response.serialize(&mut buffer); + buffer + } + }; +} + +impl CheatEngineConnection { + pub async fn run(self, handler: impl FullHandler) { + let mut buf = [0; 1024]; + + // In a loop, read data from the socket and write the data back. + let mut socket = self.0; loop { - let (mut socket, addr) = listener.accept().await?; - info!("New socket opened: {}", addr); - - tokio::spawn(async move { - let mut buf = [0; 1024]; - - // In a loop, read data from the socket and write the data back. - loop { - let n = match socket.read(&mut buf).await { - // socket closed - Ok(n) if n == 0 => return, - Ok(n) => n, - Err(e) => { - error!("failed to read from socket; err = {:?}", e); - return; - } - }; - - let request = &buf[..n]; - debug!("incoming data({}): {:?}", request[0], request); - - // Write the data back - let response = Self::handle(request); - if let Err(e) = socket.write_all(&response[..]).await { - error!("failed to write to socket; err = {:?}", e); - return; - } + let n = match socket.read(&mut buf).await { + // socket closed + Ok(n) if n == 0 => return, + Ok(n) => n, + Err(e) => { + error!("failed to read from socket; err = {:?}", e); + return; } - }); + }; + + let request = &buf[..n]; + debug!("incoming data({}): {:?}", request[0], request); + + // Write the data back + let response = Self::handle(request, &handler); + if let Err(e) = socket.write_all(&response[..]).await { + error!("failed to write to socket; err = {:?}", e); + return; + } } } - fn handle(mut bytes: &[u8]) -> Vec { - match bytes.get_u8() { - CreateToolHelp32SnapshotRequest::ID => Self::handle_packet::(&mut bytes), - Process32FirstRequest::ID => Self::handle_packet::(&mut bytes), - _ => todo!(), - } - } - - fn handle_packet(mut bytes: &[u8]) -> Vec { - let mut buffer = Vec::new(); - - let request = TReq::read(&mut bytes); - - debug!("Received item {:?}", request); - let response = request.process(); - debug!("... responding with {:?}", response); - response.serialize(&mut buffer); - - buffer + fn handle(mut bytes: &[u8], handler: &impl FullHandler) -> Vec { + gen_request_dispatch!( + bytes for handler, + CreateToolHelp32SnapshotRequest, + Process32FirstRequest + ) } }