From 02e189a02942070490168470d2e24e05abaabd64 Mon Sep 17 00:00:00 2001
From: vvb2060 <vvb2060@gmail.com>
Date: Sat, 6 Jan 2024 04:00:37 +0800
Subject: [PATCH] core: add safe mode config to db

---
 native/src/core/bootstages.cpp | 30 +++++++++++++++++++-----------
 native/src/core/daemon.rs      |  2 +-
 native/src/core/db.cpp         | 11 +++++++++++
 native/src/core/deny/utils.cpp | 12 ++----------
 native/src/core/include/db.hpp |  7 +++++--
 5 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/native/src/core/bootstages.cpp b/native/src/core/bootstages.cpp
index ba77f912e..a70ca0196 100644
--- a/native/src/core/bootstages.cpp
+++ b/native/src/core/bootstages.cpp
@@ -121,18 +121,22 @@ static bool check_key_combo() {
     return true;
 }
 
+static bool check_safe_mode() {
+    int safe_mode;
+    db_settings dbs;
+    get_db_settings(dbs, SAFEMODE_CONFIG);
+    safe_mode = dbs[SAFEMODE_CONFIG];
+
+    set_db_settings(SAFEMODE_CONFIG, safe_mode + 1);
+
+    return safe_mode >= 2 || get_prop("persist.sys.safemode", true) == "1" ||
+           get_prop("ro.sys.safemode") == "1" || check_key_combo();
+}
+
 /***********************
  * Boot Stage Handlers *
  ***********************/
 
-static void disable_zygisk() {
-    char sql[64];
-    sprintf(sql, "REPLACE INTO settings (key,value) VALUES('%s',%d)",
-            DB_SETTING_KEYS[ZYGISK_CONFIG], false);
-    char *err = db_exec(sql);
-    db_err(err);
-}
-
 bool MagiskD::post_fs_data() const {
     as_rust().setup_logfile();
 
@@ -148,21 +152,23 @@ bool MagiskD::post_fs_data() const {
             xmkdir(SECURE_DIR, 0700);
         } else {
             LOGE(SECURE_DIR " is not present, abort\n");
+            safe_mode = true;
             return safe_mode;
         }
     }
 
     if (!magisk_env()) {
         LOGE("* Magisk environment incomplete, abort\n");
+        safe_mode = true;
         return safe_mode;
     }
 
-    if (get_prop("persist.sys.safemode", true) == "1" ||
-        get_prop("ro.sys.safemode") == "1" || check_key_combo()) {
+    if (check_safe_mode()) {
+        LOGI("* Safe mode triggered\n");
         safe_mode = true;
         // Disable all modules and zygisk so next boot will be clean
         disable_modules();
-        disable_zygisk();
+        set_db_settings(ZYGISK_CONFIG, false);
         return safe_mode;
     }
 
@@ -191,6 +197,8 @@ void MagiskD::boot_complete() const {
 
     LOGI("** boot-complete triggered\n");
 
+    set_db_settings(SAFEMODE_CONFIG, 0);
+
     // At this point it's safe to create the folder
     if (access(SECURE_DIR, F_OK) != 0)
         xmkdir(SECURE_DIR, 0700);
diff --git a/native/src/core/daemon.rs b/native/src/core/daemon.rs
index 7d3f29c8c..5c0ef2b0f 100644
--- a/native/src/core/daemon.rs
+++ b/native/src/core/daemon.rs
@@ -81,7 +81,7 @@ impl MagiskD {
             }
             RequestCode::BOOT_COMPLETE => {
                 unsafe { libc::close(client) };
-                if !state.contains(BootState::SafeMode) {
+                if state.contains(BootState::PostFsDataDone) {
                     state.set(BootState::BootComplete);
                     self.as_cxx().boot_complete()
                 }
diff --git a/native/src/core/db.cpp b/native/src/core/db.cpp
index 8a8b669f5..bffbe0981 100644
--- a/native/src/core/db.cpp
+++ b/native/src/core/db.cpp
@@ -117,6 +117,7 @@ db_settings::db_settings() {
     data[SU_MNT_NS] = NAMESPACE_MODE_REQUESTER;
     data[DENYLIST_CONFIG] = false;
     data[ZYGISK_CONFIG] = MagiskD::get()->is_emulator();
+    data[SAFEMODE_CONFIG] = false;
 }
 
 int db_settings::get_idx(string_view key) const {
@@ -341,6 +342,16 @@ int get_db_settings(db_settings &cfg, int key) {
     return 0;
 }
 
+int set_db_settings(int key, int value) {
+    char *err;
+    char sql[128];
+    ssprintf(sql, sizeof(sql), "INSERT OR REPLACE INTO settings VALUES ('%s', %d)",
+             DB_SETTING_KEYS[key], value);
+    err = db_exec(sql);
+    db_err_cmd(err, return 1)
+    return 0;
+}
+
 int get_db_strings(db_strings &str, int key) {
     char *err = nullptr;
     auto string_cb = [&](db_row &row) -> bool {
diff --git a/native/src/core/deny/utils.cpp b/native/src/core/deny/utils.cpp
index dfc2bc3c4..8268170cc 100644
--- a/native/src/core/deny/utils.cpp
+++ b/native/src/core/deny/utils.cpp
@@ -324,14 +324,6 @@ void ls_list(int client) {
     close(client);
 }
 
-static void update_deny_config() {
-    char sql[64];
-    sprintf(sql, "REPLACE INTO settings (key,value) VALUES('%s',%d)",
-        DB_SETTING_KEYS[DENYLIST_CONFIG], denylist_enforced.load());
-    char *err = db_exec(sql);
-    db_err(err);
-}
-
 int enable_deny() {
     if (denylist_enforced) {
         return DenyResponse::OK;
@@ -368,7 +360,7 @@ int enable_deny() {
         }
     }
 
-    update_deny_config();
+    set_db_settings(DENYLIST_CONFIG, true);
     return DenyResponse::OK;
 }
 
@@ -376,7 +368,7 @@ int disable_deny() {
     if (denylist_enforced.exchange(false)) {
         LOGI("* Disable DenyList\n");
     }
-    update_deny_config();
+    set_db_settings(DENYLIST_CONFIG, false);
     return DenyResponse::OK;
 }
 
diff --git a/native/src/core/include/db.hpp b/native/src/core/include/db.hpp
index b434b3148..9fba8d13a 100644
--- a/native/src/core/include/db.hpp
+++ b/native/src/core/include/db.hpp
@@ -39,7 +39,8 @@ constexpr const char *DB_SETTING_KEYS[] = {
     "multiuser_mode",
     "mnt_ns",
     "denylist",
-    "zygisk"
+    "zygisk",
+    "safemode",
 };
 
 // Settings key indices
@@ -48,7 +49,8 @@ enum {
     SU_MULTIUSER_MODE,
     SU_MNT_NS,
     DENYLIST_CONFIG,
-    ZYGISK_CONFIG
+    ZYGISK_CONFIG,
+    SAFEMODE_CONFIG,
 };
 
 // Values for root_access
@@ -124,6 +126,7 @@ using db_row = std::map<std::string_view, std::string_view>;
 using db_row_cb = std::function<bool(db_row&)>;
 
 int get_db_settings(db_settings &cfg, int key = -1);
+int set_db_settings(int key, int value);
 int get_db_strings(db_strings &str, int key = -1);
 void rm_db_strings(int key);
 void exec_sql(int client);