From 22254e19dcca192fc7f54fbbcdaeaaec64df89f5 Mon Sep 17 00:00:00 2001
From: blu-dev <the.blu.dev@gmail.com>
Date: Sun, 20 Mar 2022 09:36:56 -0700
Subject: [PATCH] safely check module name and add option to request BindNow

---
 .cargo/config               |  2 --
 Xargo.toml                  | 22 ------------------
 aarch64-skyline-switch.json | 45 -------------------------------------
 rust-toolchain              |  1 -
 src/lib.rs                  | 20 +++++++++++++----
 5 files changed, 16 insertions(+), 74 deletions(-)
 delete mode 100644 .cargo/config
 delete mode 100644 Xargo.toml
 delete mode 100644 aarch64-skyline-switch.json
 delete mode 100644 rust-toolchain

diff --git a/.cargo/config b/.cargo/config
deleted file mode 100644
index 49f3dc5..0000000
--- a/.cargo/config
+++ /dev/null
@@ -1,2 +0,0 @@
-[build]
-target = "aarch64-skyline-switch"
diff --git a/Xargo.toml b/Xargo.toml
deleted file mode 100644
index 4651919..0000000
--- a/Xargo.toml
+++ /dev/null
@@ -1,22 +0,0 @@
-[package]
-rust-src = "../rust-std-skyline-squashed/src"
-
-[dependencies.core]
-path = "../rust-std-skyline-squashed/src/libcore"
-#git = "https://github.com/jam1garner/rust-std-skyline-squashed.git"
-stage = 0
-
-[dependencies.alloc]
-path = "../rust-std-skyline-squashed/src/liballoc"
-#git = "https://github.com/jam1garner/rust-std-skyline-squashed.git"
-stage = 0
-
-[dependencies.std]
-path = "../rust-std-skyline-squashed/src/libstd"
-#git = "https://github.com/jam1garner/rust-std-skyline-squashed.git"
-stage = 1
-
-#[patch.crates-io]
-#rustc-std-workspace-core = { git = 'https://github.com/jam1garner/rust-std-skyline-squashed.git', branch = 'master' }
-#rustc-std-workspace-alloc = { git = 'https://github.com/jam1garner/rust-std-skyline-squashed.git', branch = 'master' }
-#rustc-std-workspace-std = { git = 'https://github.com/jam1garner/rust-std-skyline-squashed.git', branch = 'master' }
diff --git a/aarch64-skyline-switch.json b/aarch64-skyline-switch.json
deleted file mode 100644
index b8cb884..0000000
--- a/aarch64-skyline-switch.json
+++ /dev/null
@@ -1,45 +0,0 @@
-{
-  "abi-blacklist": [
-    "stdcall",
-    "fastcall",
-    "vectorcall",
-    "thiscall",
-    "win64",
-    "sysv64"
-  ],
-  "arch": "aarch64",
-  "crt-static-default": false,
-  "crt-static-respected": false,
-  "data-layout": "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128",
-  "dynamic-linking": true,
-  "dynamic-linking-available": true,
-  "executables": true,
-  "has-elf-tls": false,
-  "has-rpath": false,
-  "linker": "rust-lld",
-  "linker-flavor": "ld.lld",
-  "llvm-target": "aarch64-unknown-none",
-  "max-atomic-width": 128,
-  "os": "switch",
-  "panic-strategy": "abort",
-  "position-independent-executables": true,
-  "pre-link-args": {
-    "ld.lld": [
-      "-Tlink.ld",
-      "-init=__custom_init",
-      "-fini=__custom_fini",
-      "--export-dynamic"
-    ]
-  },
-  "post-link-args": {
-    "ld.lld": [
-      "--no-gc-sections",
-      "--eh-frame-hdr"
-    ]
-  },
-  "relro-level": "off",
-  "target-c-int-width": "32",
-  "target-endian": "little",
-  "target-pointer-width": "64",
-  "vendor": "roblabla"
-}
diff --git a/rust-toolchain b/rust-toolchain
deleted file mode 100644
index 66822d4..0000000
--- a/rust-toolchain
+++ /dev/null
@@ -1 +0,0 @@
-nightly-2020-04-10
diff --git a/src/lib.rs b/src/lib.rs
index 0ce47d7..1d7ae92 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,9 +1,11 @@
 #![feature(proc_macro_hygiene)]
 
+use std::sync::atomic::{AtomicBool, Ordering};
+
 use skyline::{hook, install_hooks};
 use skyline::nn::ro;
 use skyline::libc::{c_void, c_int, size_t};
-use skyline::from_c_str;
+use skyline::try_from_c_str;
 
 use parking_lot::Mutex;
 
@@ -14,17 +16,22 @@ type Callback = fn(&NroInfo);
 static LOAD_HOOKS: Mutex<Vec<Callback>> = Mutex::new(Vec::new());
 static UNLOAD_HOOKS: Mutex<Vec<Callback>> = Mutex::new(Vec::new());
 
+static SHOULD_BIND_NOW: AtomicBool = AtomicBool::new(false);
+
 #[hook(replace = ro::LoadModule)]
 pub fn handle_load_module(
     out_module: *mut ro::Module, 
     image: *const c_void, 
     buffer: *mut c_void, 
     buffer_size: size_t, 
-    flag: c_int
+    mut flag: c_int
 ) -> c_int {
+    if SHOULD_BIND_NOW.load(Ordering::SeqCst) {
+        flag = 1;
+    }
     let ret = original!()(out_module, image, buffer, buffer_size, flag);
 
-    let name = unsafe { from_c_str(&(*out_module).Name as *const u8) };
+    let name = unsafe { try_from_c_str(&(*out_module).Name as *const u8).unwrap_or_else(|_| String::from("unknown")) };
     println!("[NRO hook] Loaded {}. BindFlag: {}", name, match flag {
         1 => "Now",
         2 => "Lazy",
@@ -42,7 +49,7 @@ pub fn handle_load_module(
 pub fn handle_unload_module(in_module: *mut ro::Module) -> c_int {
     let ret = original!()(in_module);
 
-    let name = unsafe { from_c_str(&(*in_module).Name as *const u8) };
+    let name = unsafe { try_from_c_str(&(*in_module).Name as *const u8).unwrap_or_else(|_| String::from("unknown")) };
     println!("[NRO hook] Unloaded {}.", name);
     let nro_info = NroInfo::new(&name, unsafe { &mut *in_module });
     for hook in UNLOAD_HOOKS.lock().iter() {
@@ -72,3 +79,8 @@ pub extern "Rust" fn add_nro_unload_hook(callback: Callback) {
 
     hooks.push(callback);
 }
+
+#[no_mangle]
+pub extern "Rust" fn nro_request_bind_now() {
+    SHOULD_BIND_NOW.store(true, Ordering::SeqCst);
+}
\ No newline at end of file