From 5e43e4a71b5231cb8f822651e23228bca34ea87a Mon Sep 17 00:00:00 2001
From: JingMatrix <jingmatrix@gmail.com>
Date: Thu, 12 Dec 2024 15:57:41 +0100
Subject: [PATCH] improve: hook pthread_attr_setstacksize

Relying on dlclose to unload libzygisk.so will block us to clean its trace in the solist.
This commit allows us to unmap libzygisk.so without using dlclose.
To call munmap, we use the function pthread_attr_setstacksize instead of pthread_attr_destroy, so that tail-call can still be applied here since it has the same signature as munmap.
---
 loader/src/injector/entry.cpp  | 11 ++++++-----
 loader/src/injector/hook.cpp   | 18 ++++++++++--------
 loader/src/injector/zygisk.hpp |  5 ++---
 loader/src/ptracer/ptracer.cpp | 18 ++++++++++++++++--
 4 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/loader/src/injector/entry.cpp b/loader/src/injector/entry.cpp
index 1cabc27..40a71b6 100644
--- a/loader/src/injector/entry.cpp
+++ b/loader/src/injector/entry.cpp
@@ -1,16 +1,17 @@
 #include "daemon.h"
 #include "logging.h"
 #include "zygisk.hpp"
-#include "module.hpp"
 
 using namespace std;
 
-void *self_handle = nullptr;
+void *start_addr = nullptr;
+size_t block_size = 0;
 
 extern "C" [[gnu::visibility("default")]]
-void entry(void* handle, const char* path) {
+void entry(void* addr, size_t size, const char* path) {
     LOGI("Zygisk library injected, version %s", ZKSU_VERSION);
-    self_handle = handle;
+    start_addr = addr;
+    block_size = size;
     zygiskd::Init(path);
 
     if (!zygiskd::PingHeartbeat()) {
@@ -22,6 +23,6 @@ void entry(void* handle, const char* path) {
     logging::setfd(zygiskd::RequestLogcatFd());
 #endif
 
-    LOGI("Start hooking");
+    LOGI("Start hooking, call %p", hook_functions);
     hook_functions();
 }
diff --git a/loader/src/injector/hook.cpp b/loader/src/injector/hook.cpp
index a1b715c..16bb8d1 100644
--- a/loader/src/injector/hook.cpp
+++ b/loader/src/injector/hook.cpp
@@ -176,22 +176,24 @@ DCL_HOOK_FUNC(void, android_log_close) {
 
 // We cannot directly call `dlclose` to unload ourselves, otherwise when `dlclose` returns,
 // it will return to our code which has been unmapped, causing segmentation fault.
-// Instead, we hook `pthread_attr_destroy` which will be called when VM daemon threads start.
-DCL_HOOK_FUNC(int, pthread_attr_destroy, void *target) {
-    int res = old_pthread_attr_destroy((pthread_attr_t *)target);
+// Instead, we hook `pthread_attr_setstacksize` which will be called when VM daemon threads start.
+DCL_HOOK_FUNC(int, pthread_attr_setstacksize, void *target, size_t size) {
+    int res = old_pthread_attr_setstacksize((pthread_attr_t *)target, size);
+    LOGV("Call pthread_attr_setstacksize in [tid, pid]: %d, %d", gettid(), getpid());
 
     // Only perform unloading on the main thread
     if (gettid() != getpid())
         return res;
 
-    LOGV("pthread_attr_destroy");
+    LOGV("Clean zygisk reminders");
     if (should_unmap_zygisk) {
         unhook_functions();
         if (should_unmap_zygisk) {
-            // Because both `pthread_attr_destroy` and `dlclose` have the same function signature,
+            // Because both `pthread_attr_setstacksize` and `dlclose` have the same function signature,
             // we can use `musttail` to let the compiler reuse our stack frame and thus
-            // `dlclose` will directly return to the caller of `pthread_attr_destroy`.
-            [[clang::musttail]] return dlclose(self_handle);
+            // `dlclose` will directly return to the caller of `pthread_attr_setstacksize`.
+            LOGV("Unmap libzygisk.so");
+            [[clang::musttail]] return munmap(start_addr, block_size);
         }
     }
 
@@ -858,7 +860,7 @@ static void hook_unloader() {
         }
     }
 
-    PLT_HOOK_REGISTER(art_dev, art_inode, pthread_attr_destroy);
+    PLT_HOOK_REGISTER(art_dev, art_inode, pthread_attr_setstacksize);
     hook_commit();
 }
 
diff --git a/loader/src/injector/zygisk.hpp b/loader/src/injector/zygisk.hpp
index 5c06f8b..636367b 100644
--- a/loader/src/injector/zygisk.hpp
+++ b/loader/src/injector/zygisk.hpp
@@ -1,10 +1,9 @@
 #pragma once
 
-#include <stdint.h>
 #include <jni.h>
-#include <vector>
 
-extern void *self_handle;
+extern void *start_addr;
+extern size_t block_size;
 
 void hook_functions();
 
diff --git a/loader/src/ptracer/ptracer.cpp b/loader/src/ptracer/ptracer.cpp
index 9182ddd..f790f98 100644
--- a/loader/src/ptracer/ptracer.cpp
+++ b/loader/src/ptracer/ptracer.cpp
@@ -222,10 +222,24 @@ bool inject_on_main(int pid, const char *lib_path) {
       return false;
     }
 
-    /* call injector entry(handle, path) */
+    /* record the address range of libzygisk.so */
+    map = MapInfo::Scan(std::to_string(pid));
+    void *start_addr = nullptr;
+    size_t block_size = 0;
+    for (auto &info : map) {
+        if (strstr(info.path.c_str(), "libzygisk.so")) {
+            void *addr = (void *)info.start;
+            if (start_addr == nullptr) start_addr = addr;
+            size_t size = info.end - info.start;
+            block_size += size;
+            LOGD("found block %s: [%p-%p] with size %zu", info.path.c_str(), addr, (void *)info.end, size);
+        }
+    }
 
+    /* call injector entry(start_addr, block_size, path) */
     args.clear();
-    args.push_back(remote_handle);
+    args.push_back((uintptr_t) start_addr);
+    args.push_back(block_size);
     str = push_string(pid, regs, zygiskd::GetTmpPath().c_str());
     args.push_back((long) str);