From 40ac7d3dcd59b14a1318594e1ea7357655f2fe47 Mon Sep 17 00:00:00 2001
From: xavo95 <xavo95@xeondev.com>
Date: Fri, 1 Nov 2024 20:30:07 +0100
Subject: [PATCH] Remove /index.json redirect and allow customizable config
 server by env

---
 Cargo.lock      | 12 ++++++++++--
 Cargo.toml      |  2 ++
 src/lib.rs      | 27 +++++++++++++++++++++------
 src/replacer.rs | 29 +++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 8 deletions(-)
 create mode 100644 src/replacer.rs

diff --git a/Cargo.lock b/Cargo.lock
index 58e322f..f96088b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -22,6 +22,8 @@ name = "camellya-patch"
 version = "0.1.0"
 dependencies = [
  "ilhook",
+ "regex",
+ "widestring",
  "windows",
 ]
 
@@ -87,9 +89,9 @@ dependencies = [
 
 [[package]]
 name = "regex"
-version = "1.11.0"
+version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8"
+checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -151,6 +153,12 @@ version = "1.0.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe"
 
+[[package]]
+name = "widestring"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7219d36b6eac893fa81e84ebe06485e7dcbb616177469b142df14f1f4deb1311"
+
 [[package]]
 name = "windows"
 version = "0.58.0"
diff --git a/Cargo.toml b/Cargo.toml
index 92cd947..0ba0345 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,6 +16,8 @@ enable-sdk = []
 
 [dependencies]
 ilhook = "2.1.1"
+regex = "1.11.1"
+widestring = "1.1.0"
 windows = { version = "0.58.0", features = [
 	"Win32_Foundation",
 	"Win32_System_LibraryLoader",
diff --git a/src/lib.rs b/src/lib.rs
index b5b5ecc..1be75d3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,8 +1,9 @@
+use std::sync::OnceLock;
 use std::thread;
 use std::time::Duration;
 
 use ilhook::x64::Registers;
-use windows::core::{PCSTR, PCWSTR, w};
+use windows::core::{PCSTR, PCWSTR};
 use windows::Win32::Foundation::HINSTANCE;
 use windows::Win32::System::Console;
 use windows::Win32::System::LibraryLoader::GetModuleHandleA;
@@ -11,8 +12,13 @@ use windows::Win32::System::SystemServices::DLL_PROCESS_ATTACH;
 use interceptor::Interceptor;
 use offsets::CONFIG;
 
+use crate::replacer::{GenericReplacer, Replacer};
+
 mod interceptor;
 mod offsets;
+mod replacer;
+
+static CFG_SERVER_REPLACER: OnceLock<GenericReplacer> = OnceLock::new();
 
 fn thread_func() {
     unsafe { Console::AllocConsole() }.unwrap();
@@ -64,13 +70,22 @@ fn thread_func() {
 unsafe extern "win64" fn on_kurohttp_get(reg: *mut Registers, _: usize) {
     let wstr = *((*reg).rcx as *const usize) as *mut u16;
     let url = PCWSTR::from_raw(wstr).to_string().unwrap();
-
     println!("HTTP GET: {url}");
-    if url.ends_with("/index.json") {
-        println!("index.json requested, redirecting");
-        let new_wstr = w!("http://127.0.0.1:10001/index.json");
+    
+    let replacer = CFG_SERVER_REPLACER.get_or_init(|| {
+        GenericReplacer {
+            regex: regex::Regex::new(r#"^(?:https|http)://.*/([a-zA-Z0-9]{32}/index\.json)$"#).unwrap(),
+            replacement: std::env::var("CFG_SERVER_URL").unwrap_or("127.0.0.1:10001".to_string()),
+            scheme: std::env::var("CFG_SERVER_SCHEME").unwrap_or("http".to_string()),
+        }
+    });
+    
+    if let Some(result) = replacer.replace(url.as_str()) {
+        println!("Redirecting to: {result}");
+        let new_url = widestring::U16String::from_str(result.as_str());
+        let new_wstr = PCWSTR::from_raw(new_url.as_ptr());
         std::ptr::copy_nonoverlapping(new_wstr.as_ptr(), wstr, new_wstr.as_wide().len() + 2);
-    }
+    };
 }
 
 unsafe extern "win64" fn fpakfile_check_replacement(
diff --git a/src/replacer.rs b/src/replacer.rs
new file mode 100644
index 0000000..68a8e44
--- /dev/null
+++ b/src/replacer.rs
@@ -0,0 +1,29 @@
+pub(crate) trait Replacer {
+    fn replace(&self, original: &str) -> Option<String>;
+}
+
+pub(crate) struct GenericReplacer {
+    pub(crate) regex: regex::Regex,
+    pub(crate) replacement: String,
+    pub(crate) scheme: String,
+}
+
+impl Replacer for GenericReplacer {
+    fn replace(&self, original: &str) -> Option<String> {
+        // Prepare output array
+        let mut results: Vec<String> = vec![];
+        // Perform the capture over input
+        for (_, [path]) in self.regex.captures_iter(original).map(|c| c.extract()) {
+            results.push(format!("{}://{}/{}", self.scheme, self.replacement, path));
+        }
+        // We are supposed to only parse one entry from text
+        if 1 == results.len() {
+            return Some(results.remove(0));
+        } else if results.is_empty() {
+            println!("No valid url match found so returning original url");
+        } else {
+            println!("Invalid number of entries parsed, expected 1, obtained {:?}", results.len());
+        }
+        None
+    }
+}
\ No newline at end of file