From 369b94e7e355de562f1787b70f7b1a47aca8ef8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Richard=20Iv=C3=A1nek?= Date: Thu, 26 Sep 2024 11:07:10 +0200 Subject: [PATCH] refactor: simplify code, use format!, more idiomatic code --- Cargo.lock | 33 +++++++++++++++++++++ Cargo.toml | 2 ++ src/consts.rs | 2 +- src/jobs.rs | 81 +++++++++++++++++++++------------------------------ src/main.rs | 58 +++++++++++++++--------------------- src/opcode.rs | 2 +- src/player.rs | 57 +++++++++++++++--------------------- 7 files changed, 117 insertions(+), 118 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8c44f2a..6d633a7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -394,6 +394,12 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.3.9" @@ -533,6 +539,8 @@ dependencies = [ "regex", "reqwest", "rquickjs", + "strum", + "strum_macros", "tokio", "tokio-util", "tub", @@ -969,6 +977,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "beb461507cee2c2ff151784c52762cf4d9ff6a61f3e80968600ed24fa837fa54" +[[package]] +name = "rustversion" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "955d28af4278de8121b7ebeb796b6a45735dc01436d898801014aced2773a3d6" + [[package]] name = "ryu" version = "1.0.17" @@ -1090,6 +1104,25 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "syn" version = "2.0.60" diff --git a/Cargo.toml b/Cargo.toml index af3d3f9..2b13bfd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,8 @@ tokio-util = { version = "0.7.10", features=["futures-io", "futures-util", "code futures = "0.3.30" log = "0.4.22" env_logger = "0.11.5" +strum = "0.26.3" +strum_macros = "0.26.4" # Compilation optimizations for release builds # Increases compile time but typically produces a faster and smaller binary. Suitable for final releases but not for debug builds. diff --git a/src/consts.rs b/src/consts.rs index 33f2f4e..f0beef2 100644 --- a/src/consts.rs +++ b/src/consts.rs @@ -25,4 +25,4 @@ pub static REGEX_SIGNATURE_FUNCTION: &Lazy = pub static REGEX_HELPER_OBJ_NAME: &Lazy = regex!(";([A-Za-z0-9_\\$]{2,})\\...\\("); pub static NSIG_FUNCTION_NAME: &str = "decrypt_nsig"; -pub static SIG_FUNCTION_NAME: &str = "decrypt_sig"; +pub static _SIG_FUNCTION_NAME: &str = "decrypt_sig"; diff --git a/src/jobs.rs b/src/jobs.rs index ae7ae11..ba7ac3a 100644 --- a/src/jobs.rs +++ b/src/jobs.rs @@ -1,46 +1,27 @@ use futures::SinkExt; +use log::{debug, error}; use rquickjs::{async_with, AsyncContext, AsyncRuntime}; use std::{num::NonZeroUsize, sync::Arc, thread::available_parallelism, time::SystemTime}; -use log::{debug, error}; +use strum_macros::{Display, FromRepr}; use tokio::{runtime::Handle, sync::Mutex, task::block_in_place}; use tub::Pool; use crate::{consts::NSIG_FUNCTION_NAME, opcode::OpcodeResponse, player::fetch_update}; +#[derive(Display, FromRepr)] pub enum JobOpcode { - ForceUpdate, - DecryptNSignature, + ForceUpdate = 0, + DecryptNSignature = 1, DecryptSignature, GetSignatureTimestamp, PlayerStatus, PlayerUpdateTimestamp, - UnknownOpcode, + UnknownOpcode = 255, } -impl std::fmt::Display for JobOpcode { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::ForceUpdate => write!(f, "ForceUpdate"), - Self::DecryptNSignature => write!(f, "DecryptNSignature"), - Self::DecryptSignature => write!(f, "DecryptSignature"), - Self::GetSignatureTimestamp => write!(f, "GetSignatureTimestamp"), - Self::PlayerStatus => write!(f, "PlayerStatus"), - Self::PlayerUpdateTimestamp => write!(f, "PlayerUpdateTimestamp"), - Self::UnknownOpcode => write!(f, "UnknownOpcode"), - } - } -} impl From for JobOpcode { fn from(value: u8) -> Self { - match value { - 0x00 => Self::ForceUpdate, - 0x01 => Self::DecryptNSignature, - 0x02 => Self::DecryptSignature, - 0x03 => Self::GetSignatureTimestamp, - 0x04 => Self::PlayerStatus, - 0x05 => Self::PlayerUpdateTimestamp, - _ => Self::UnknownOpcode, - } + JobOpcode::from_repr(value as usize).unwrap_or(Self::UnknownOpcode) } } @@ -54,7 +35,22 @@ pub struct PlayerInfo { pub last_update: SystemTime, } +impl Default for PlayerInfo { + fn default() -> Self { + Self { + nsig_function_code: Default::default(), + sig_function_code: Default::default(), + sig_function_name: Default::default(), + signature_timestamp: Default::default(), + player_id: Default::default(), + has_player: Default::default(), + last_update: SystemTime::now(), + } + } +} + pub struct JavascriptInterpreter { + #[allow(dead_code)] js_runtime: AsyncRuntime, sig_context: AsyncContext, nsig_context: AsyncContext, @@ -65,17 +61,20 @@ pub struct JavascriptInterpreter { impl JavascriptInterpreter { pub fn new() -> JavascriptInterpreter { let js_runtime = AsyncRuntime::new().unwrap(); + // not ideal, but this is only done at startup let nsig_context = block_in_place(|| { Handle::current() .block_on(AsyncContext::full(&js_runtime)) .unwrap() }); + let sig_context = block_in_place(|| { Handle::current() .block_on(AsyncContext::full(&js_runtime)) .unwrap() }); + JavascriptInterpreter { js_runtime, sig_context, @@ -98,22 +97,16 @@ impl GlobalState { .get(); let mut runtime_vector: Vec> = Vec::with_capacity(number_of_runtimes); - for _n in 0..number_of_runtimes { + + for _ in 0..number_of_runtimes { runtime_vector.push(Arc::new(JavascriptInterpreter::new())); } - let runtime_pool: Pool> = Pool::from_vec(runtime_vector); + let js_runtime_pool: Pool> = Pool::from_vec(runtime_vector); + GlobalState { - player_info: Mutex::new(PlayerInfo { - nsig_function_code: Default::default(), - sig_function_code: Default::default(), - sig_function_name: Default::default(), - player_id: Default::default(), - signature_timestamp: Default::default(), - has_player: 0x00, - last_update: SystemTime::now(), - }), - js_runtime_pool: runtime_pool, + player_info: Mutex::new(PlayerInfo::default()), + js_runtime_pool, } } } @@ -183,11 +176,7 @@ pub async fn process_decrypt_n_signature( } drop(player_info); - let mut call_string: String = String::new(); - call_string += NSIG_FUNCTION_NAME; - call_string += "(\""; - call_string += &sig.replace("\"", "\\\""); - call_string += "\")"; + let call_string = format!("{NSIG_FUNCTION_NAME}(\"{}\")", sig.replace("\"", "\\\"")); let decrypted_string = match ctx.eval::(call_string.clone()) { Ok(x) => x, @@ -263,11 +252,7 @@ pub async fn process_decrypt_signature( let sig_function_name = &player_info.sig_function_name; - let mut call_string: String = String::new(); - call_string += sig_function_name; - call_string += "(\""; - call_string += &sig.replace("\"", "\\\""); - call_string += "\")"; + let call_string = format!("{sig_function_name}(\"{}\")", sig.replace("\"", "\\\"")); drop(player_info); diff --git a/src/main.rs b/src/main.rs index fad8b8f..d6c7bda 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,11 +5,12 @@ mod player; use ::futures::StreamExt; use consts::{DEFAULT_SOCK_PATH, DEFAULT_TCP_URL}; +use env_logger::Env; use jobs::{process_decrypt_n_signature, process_fetch_update, GlobalState, JobOpcode}; +use log::{debug, error, info}; use opcode::OpcodeDecoder; use player::fetch_update; use std::{env::args, sync::Arc}; -use env_logger::Env; use tokio::{ fs::remove_file, io::{AsyncReadExt, AsyncWrite}, @@ -17,7 +18,6 @@ use tokio::{ sync::Mutex, }; use tokio_util::codec::Framed; -use log::{info, error, debug}; use crate::jobs::{ process_decrypt_signature, process_get_signature_timestamp, process_player_status, @@ -30,7 +30,7 @@ macro_rules! loop_main { match fetch_update($s.clone()).await { Ok(()) => info!("Successfully fetched player"), Err(x) => { - error!("Error occured while trying to fetch the player: {:?}", x); + error!("Error occurred while trying to fetch the player: {:?}", x); } } loop { @@ -48,19 +48,14 @@ async fn main() { env_logger::Builder::from_env(Env::default().default_filter_or("info")).init(); let args: Vec = args().collect(); - let socket_url: &str = match args.get(1) { - Some(stringref) => stringref, - None => DEFAULT_SOCK_PATH, - }; + let socket_url: &str = args.get(1).map(String::as_ref).unwrap_or(DEFAULT_SOCK_PATH); // have to please rust let state: Arc = Arc::new(GlobalState::new()); if socket_url == "--tcp" { - let socket_tcp_url: &str = match args.get(2) { - Some(stringref) => stringref, - None => DEFAULT_TCP_URL, - }; + let socket_tcp_url = args.get(2).map(String::as_ref).unwrap_or(DEFAULT_TCP_URL); + let tcp_socket = match TcpListener::bind(socket_tcp_url).await { Ok(x) => x, Err(x) => { @@ -68,27 +63,29 @@ async fn main() { return; } }; + loop_main!(tcp_socket, state); } else if socket_url == "--test" { - // TODO: test the API aswell, this only tests the player script extractor + // TODO: test the API as well, this only tests the player script extractor info!("Fetching player"); - match fetch_update(state.clone()).await { - Ok(()) => std::process::exit(0), - Err(_x) => std::process::exit(-1), - } + + std::process::exit(match fetch_update(state.clone()).await { + Ok(_) => 0, + Err(_) => -1, + }); } else { let unix_socket = match UnixListener::bind(socket_url) { Ok(x) => x, + Err(x) if x.kind() == std::io::ErrorKind::AddrInUse => { + let _ = remove_file(socket_url).await; + UnixListener::bind(socket_url).unwrap() + } Err(x) => { - if x.kind() == std::io::ErrorKind::AddrInUse { - remove_file(socket_url).await; - UnixListener::bind(socket_url).unwrap() - } else { - error!("Error occurred while trying to bind: {}", x); - return; - } + error!("Error occurred while trying to bind: {}", x); + return; } }; + loop_main!(unix_socket, state); } } @@ -108,18 +105,17 @@ where Ok(opcode) => { debug!("Received job: {}", opcode.opcode); + let cloned_state = state.clone(); + let cloned_sink = arc_sink.clone(); + match opcode.opcode { JobOpcode::ForceUpdate => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_fetch_update(cloned_state, cloned_sink, opcode.request_id) .await; }); } JobOpcode::DecryptNSignature => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_decrypt_n_signature( cloned_state, @@ -131,8 +127,6 @@ where }); } JobOpcode::DecryptSignature => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_decrypt_signature( cloned_state, @@ -144,8 +138,6 @@ where }); } JobOpcode::GetSignatureTimestamp => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_get_signature_timestamp( cloned_state, @@ -156,16 +148,12 @@ where }); } JobOpcode::PlayerStatus => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_player_status(cloned_state, cloned_sink, opcode.request_id) .await; }); } JobOpcode::PlayerUpdateTimestamp => { - let cloned_state = state.clone(); - let cloned_sink = arc_sink.clone(); tokio::spawn(async move { process_player_update_timestamp( cloned_state, diff --git a/src/opcode.rs b/src/opcode.rs index d747f0e..d0c1d8f 100644 --- a/src/opcode.rs +++ b/src/opcode.rs @@ -1,5 +1,5 @@ -use std::io::ErrorKind; use log::debug; +use std::io::ErrorKind; use tokio_util::{ bytes::{Buf, BufMut}, codec::{Decoder, Encoder}, diff --git a/src/player.rs b/src/player.rs index 23333b3..3be56b0 100644 --- a/src/player.rs +++ b/src/player.rs @@ -1,6 +1,6 @@ -use std::{sync::Arc, time::SystemTime}; use log::{debug, error, info, warn}; use regex::Regex; +use std::{sync::Arc, time::SystemTime}; use crate::{ consts::{ @@ -64,19 +64,20 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat let mut nsig_function_array_opt = None; // Extract nsig function array code for (index, nsig_function_array_str) in NSIG_FUNCTION_ARRAYS.iter().enumerate() { - let nsig_function_array_regex = Regex::new(&nsig_function_array_str).unwrap(); + let nsig_function_array_regex = Regex::new(nsig_function_array_str).unwrap(); nsig_function_array_opt = match nsig_function_array_regex.captures(&player_javascript) { None => { - warn!("nsig function array did not work: {}", nsig_function_array_str); + warn!( + "nsig function array did not work: {}", + nsig_function_array_str + ); if index == NSIG_FUNCTION_ARRAYS.len() { error!("!!ERROR!! nsig function array unable to be extracted"); return Err(FetchUpdateStatus::NsigRegexCompileFailed); } continue; } - Some(i) => { - Some(i) - } + Some(i) => Some(i), }; break; } @@ -90,10 +91,10 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat .parse::() .unwrap(); - let mut nsig_array_context_regex: String = String::new(); - nsig_array_context_regex += "var "; - nsig_array_context_regex += &nsig_array_name.replace("$", "\\$"); - nsig_array_context_regex += "\\s*=\\s*\\[(.+?)][;,]"; + let nsig_array_context_regex: String = format!( + "var {name}\\s*=\\s*\\[(.+?)][;,]", + name = nsig_array_name.replace("$", "\\$") + ); let nsig_array_context = match Regex::new(&nsig_array_context_regex) { Ok(x) => x, @@ -115,20 +116,19 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat let nsig_function_name = array_values.get(nsig_array_value).unwrap(); - let mut nsig_function_code = String::new(); - nsig_function_code += "function "; - nsig_function_code += NSIG_FUNCTION_NAME; + let mut nsig_function_code = format!("function {NSIG_FUNCTION_NAME}"); // Extract nsig function code for (index, ending) in NSIG_FUNCTION_ENDINGS.iter().enumerate() { - let mut nsig_function_code_regex_str: String = String::new(); - nsig_function_code_regex_str += &nsig_function_name.replace("$", "\\$"); - nsig_function_code_regex_str += ending; + let nsig_function_code_regex_str = + format!("{}{ending}", nsig_function_name.replace("$", "\\$")); let nsig_function_code_regex = Regex::new(&nsig_function_code_regex_str).unwrap(); + nsig_function_code += match nsig_function_code_regex.captures(&player_javascript) { None => { warn!("nsig function ending did not work: {}", ending); + if index == NSIG_FUNCTION_ENDINGS.len() { error!("!!ERROR!! nsig function unable to be extracted"); return Err(FetchUpdateStatus::NsigRegexCompileFailed); @@ -136,9 +136,7 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat continue; } - Some(i) => { - i.get(1).unwrap().as_str() - } + Some(i) => i.get(1).unwrap().as_str(), }; debug!("got nsig fn code: {}", nsig_function_code); break; @@ -152,9 +150,10 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat .unwrap() .as_str(); - let mut sig_function_body_regex_str: String = String::new(); - sig_function_body_regex_str += &sig_function_name.replace("$", "\\$"); - sig_function_body_regex_str += "=function\\([a-zA-Z0-9_]+\\)\\{.+?\\}"; + let sig_function_body_regex_str = format!( + "{name}=function\\([a-zA-Z0-9_]+\\)\\{{.+?\\}}", + name = sig_function_name.replace("$", "\\$") + ); let sig_function_body_regex = Regex::new(&sig_function_body_regex_str).unwrap(); @@ -173,10 +172,8 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat .unwrap() .as_str(); - let mut helper_object_body_regex_str = String::new(); - helper_object_body_regex_str += "(var "; - helper_object_body_regex_str += helper_object_name; - helper_object_body_regex_str += "=\\{(?:.|\\n)+?\\}\\};)"; + let helper_object_body_regex_str = + format!("(var {helper_object_name}=\\{{(?:.|\\n)+?\\}}\\}};)"); let helper_object_body_regex = Regex::new(&helper_object_body_regex_str).unwrap(); let helper_object_body = helper_object_body_regex @@ -186,13 +183,7 @@ pub async fn fetch_update(state: Arc) -> Result<(), FetchUpdateStat .unwrap() .as_str(); - let mut sig_code = String::new(); - sig_code += "var "; - sig_code += sig_function_name; - sig_code += ";"; - - sig_code += helper_object_body; - sig_code += sig_function_body; + let sig_code = format!("var {sig_function_name};{helper_object_body}{sig_function_body}"); info!("sig code: {}", sig_code);