From 8042e303ace4ee967f6a049a9fb40f26dfc21907 Mon Sep 17 00:00:00 2001
From: Adamaq01 <adamthibert01@gmail.com>
Date: Sun, 16 Jul 2023 08:30:16 +0200
Subject: [PATCH] fix: re-encode hijacked responses the same way as the
 originals

---
 src/cloudlink/mod.rs | 10 +++-------
 src/mikado.rs        | 37 +++++++++++++++++++++++++------------
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/cloudlink/mod.rs b/src/cloudlink/mod.rs
index 1d4c302..1094f58 100644
--- a/src/cloudlink/mod.rs
+++ b/src/cloudlink/mod.rs
@@ -5,7 +5,7 @@ use crate::{helpers, TACHI_PBS_URL};
 use anyhow::Result;
 use dynfmt::Format;
 use ext::HashMapExt;
-use kbinxml::{CompressionType, EncodingType, Node, Options, Value, ValueArray};
+use kbinxml::{Node, Value, ValueArray};
 use log::info;
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
@@ -20,7 +20,7 @@ fn build_response_base(scores: Vec<Node>) -> Node {
     )
 }
 
-pub fn process_pbs(user: &str, music: &Node, encoding: EncodingType) -> Result<Vec<u8>> {
+pub fn process_pbs(user: &str, music: &Node) -> Result<Node> {
     let url = dynfmt::SimpleCurlyFormat.format(TACHI_PBS_URL.as_str(), [user])?;
     let response: serde_json::Value = helpers::request_tachi("GET", url, None::<()>)?;
     let body = response["body"].as_object().ok_or(anyhow::anyhow!(
@@ -120,11 +120,7 @@ pub fn process_pbs(user: &str, music: &Node, encoding: EncodingType) -> Result<V
     }
 
     let response = build_response_base(scores.to_properties());
-    let bytes = kbinxml::to_binary_with_options(
-        Options::new(CompressionType::Uncompressed, encoding),
-        &response,
-    )?;
     info!("Successfully injected Tachi PBs as Cloud scores");
 
-    Ok(bytes)
+    Ok(response)
 }
diff --git a/src/mikado.rs b/src/mikado.rs
index a787fff..49828c0 100644
--- a/src/mikado.rs
+++ b/src/mikado.rs
@@ -3,7 +3,7 @@ use std::sync::RwLock;
 
 use anyhow::Result;
 use bytes::Bytes;
-use kbinxml::{CompressionType, Node, Options, Value};
+use kbinxml::{CompressionType, EncodingType, Node, Options, Value};
 use log::{debug, error, info, warn};
 
 use crate::handlers::save::process_save;
@@ -132,14 +132,32 @@ pub unsafe fn property_mem_read_hook(
     }
 }
 
+fn build_response(
+    original_signature: &[u8],
+    response: Node,
+    encoding: EncodingType,
+) -> Result<Vec<u8>> {
+    if kbinxml::is_binary_xml(original_signature) {
+        let bytes = kbinxml::to_binary_with_options(
+            Options::new(CompressionType::from_byte(original_signature[1])?, encoding),
+            &response,
+        )?;
+        Ok(bytes)
+    } else {
+        let bytes = kbinxml::to_text_xml(&response)?;
+        Ok(bytes)
+    }
+}
+
 #[allow(clippy::manual_map)]
 pub unsafe fn property_mem_read_hook_wrapped(
-    response: Vec<u8>,
+    original: Vec<u8>,
     load: bool,
     load_m: bool,
     common: bool,
 ) -> Option<Result<Vec<u8>>> {
-    let (mut root, encoding) = kbinxml::from_bytes(Bytes::from(response))
+    let original_signature = original[..2].to_vec();
+    let (mut root, encoding) = kbinxml::from_bytes(Bytes::from(original))
         .and_then(|(node, encoding)| node.as_node().map(|node| (node, encoding)))
         .ok()?;
 
@@ -170,10 +188,7 @@ pub unsafe fn property_mem_read_hook_wrapped(
                     Value::String("CLOUD_LINK_ENABLE".to_string()),
                 )],
             ));
-            let response = kbinxml::to_binary_with_options(
-                Options::new(CompressionType::Uncompressed, encoding),
-                &root,
-            )?;
+            let response = build_response(&original_signature, root, encoding)?;
             COMMON.store(false, Ordering::Relaxed);
 
             Ok(response)
@@ -194,10 +209,7 @@ pub unsafe fn property_mem_read_hook_wrapped(
                 "cloud",
                 vec![Node::with_value("relation", Value::S8(1))],
             ));
-            let response = kbinxml::to_binary_with_options(
-                Options::new(CompressionType::Uncompressed, encoding),
-                &root,
-            )?;
+            let response = build_response(&original_signature, root, encoding)?;
             LOAD.store(false, Ordering::Relaxed);
 
             Ok(response)
@@ -205,7 +217,8 @@ pub unsafe fn property_mem_read_hook_wrapped(
     } else if let Some(music) = load_m.then(|| root.pointer(&["game", "music"])).flatten() {
         Some((|| {
             let user = USER.load(Ordering::SeqCst).to_string();
-            let response = crate::cloudlink::process_pbs(user.as_str(), music, encoding)?;
+            let response = crate::cloudlink::process_pbs(user.as_str(), music)?;
+            let response = build_response(&original_signature, response, encoding)?;
             LOAD_M.store(false, Ordering::Relaxed);
 
             Ok(response)