Made reboot to payload safer

Makefile

@@ -22,7 +22,7 @@ DATA		:=	data
 INCLUDES	:=	include lib/zipper/include /lib/borealis/library/include/borealis/extern/nlohmann
 APP_TITLE	:=	All-in-One Switch Updater
 APP_AUTHOR	:=	HamletDuFromage
-APP_VERSION :=  2.6.0
+APP_VERSION :=  2.6.1
 TARGET		:=	$(notdir $(CURDIR))
 
 ROMFS				:=	resources

include/ams_bpc.h

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2018-2020 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#pragma once
+#include <switch.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+Result amsBpcInitialize();
+void   amsBpcExit();
+Service *amsBpcGetServiceSession(void);
+
+Result amsBpcSetRebootPayload(const void *src, size_t src_size);
+
+#ifdef __cplusplus
+}
+#endif

include/reboot_payload.h

@@ -10,7 +10,7 @@ extern "C" {
 #include <stdbool.h>
 #include <switch.h>
 
-int reboot_to_payload(const char* path);
+int reboot_to_payload(const char* path, bool legacy);
 
 #ifdef __cplusplus
 }

include/service_guard.h

@@ -0,0 +1,56 @@
+#pragma once
+#include <switch/types.h>
+#include <switch/result.h>
+#include <switch/kernel/mutex.h>
+#include <switch/sf/service.h>
+#include <switch/services/sm.h>
+
+typedef struct ServiceGuard {
+    Mutex mutex;
+    u32 refCount;
+} ServiceGuard;
+
+NX_INLINE bool serviceGuardBeginInit(ServiceGuard* g)
+{
+    mutexLock(&g->mutex);
+    return (g->refCount++) == 0;
+}
+
+NX_INLINE Result serviceGuardEndInit(ServiceGuard* g, Result rc, void (*cleanupFunc)(void))
+{
+    if (R_FAILED(rc)) {
+        cleanupFunc();
+        --g->refCount;
+    }
+    mutexUnlock(&g->mutex);
+    return rc;
+}
+
+NX_INLINE void serviceGuardExit(ServiceGuard* g, void (*cleanupFunc)(void))
+{
+    mutexLock(&g->mutex);
+    if (g->refCount && (--g->refCount) == 0)
+        cleanupFunc();
+    mutexUnlock(&g->mutex);
+}
+
+#define NX_GENERATE_SERVICE_GUARD_PARAMS(name, _paramdecl, _parampass) \
+\
+static ServiceGuard g_##name##Guard; \
+NX_INLINE Result _##name##Initialize _paramdecl; \
+static void _##name##Cleanup(void); \
+\
+Result name##Initialize _paramdecl \
+{ \
+    Result rc = 0; \
+    if (serviceGuardBeginInit(&g_##name##Guard)) \
+        rc = _##name##Initialize _parampass; \
+    return serviceGuardEndInit(&g_##name##Guard, rc, _##name##Cleanup); \
+} \
+\
+void name##Exit(void) \
+{ \
+    serviceGuardExit(&g_##name##Guard, _##name##Cleanup); \
+}
+
+#define NX_GENERATE_SERVICE_GUARD(name) NX_GENERATE_SERVICE_GUARD_PARAMS(name, (void), ())

source/ams_bpc.c

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) 2018-2020 Atmosphère-NX
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms and conditions of the GNU General Public License,
+ * version 2, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+#include <switch.h>
+#include <string.h>
+#include "ams_bpc.h"
+#include "service_guard.h"
+
+static Service g_amsBpcSrv;
+
+NX_GENERATE_SERVICE_GUARD(amsBpc);
+
+Result _amsBpcInitialize(void) {
+    Handle h;
+    Result rc = svcConnectToNamedPort(&h, "bpc:ams"); /* TODO: ams:bpc */
+    if (R_SUCCEEDED(rc)) serviceCreate(&g_amsBpcSrv, h);
+    return rc;
+}
+
+void _amsBpcCleanup(void) {
+    serviceClose(&g_amsBpcSrv);
+}
+
+Service *amsBpcGetServiceSession(void) {
+    return &g_amsBpcSrv;
+}
+
+Result amsBpcSetRebootPayload(const void *src, size_t src_size) {
+    return serviceDispatch(&g_amsBpcSrv, 65001,
+        .buffer_attrs = { SfBufferAttr_In | SfBufferAttr_HipcMapAlias },
+        .buffers = { { src, src_size } },
+    );
+}

source/changelog_page.cpp

@@ -160,6 +160,9 @@ ChangelogPage::ChangelogPage() : AppletFrame(true, true)
     verTitles.push_back("v2.6.0");
     changes.push_back("\uE016 Added ability to download individual cheat codes from the GBAtemp.net archive.\n\uE016 Updated Polish translation (credits to github.com/teddy74eva).");
 
+    verTitles.push_back("v2.6.0");
+    changes.push_back("\uE016 Added safer way to reboot to payload (should prevent corruption on exFAT system, not that should be using exFAT)");
+
     for(int i = verTitles.size() -1 ; i >= 0; i--){
         listItem = new brls::ListItem(verTitles[i]);
         change = changes[i];

source/reboot_payload.c

@@ -1,4 +1,5 @@
 #include "reboot_payload.h"
+#include "ams_bpc.h"
 
 #define IRAM_PAYLOAD_MAX_SIZE   0x2F000
 #define IRAM_PAYLOAD_BASE       0x40010000
@@ -38,29 +39,36 @@ static void clear_iram(void) {
 }
 
 static void inject_payload(void) {
+    printf("injecting\n");
+    spsmInitialize();
+    smExit();
+    if (R_SUCCEEDED(amsBpcInitialize()) && R_SUCCEEDED(amsBpcSetRebootPayload(g_reboot_payload, 0x24000))) {
+        spsmShutdown(true);
+    }
+}
+
+static void inject_payload_legacy(void) {
+    printf("injecting (legacy)\n");
     clear_iram();
-    
     for (size_t i = 0; i < IRAM_PAYLOAD_MAX_SIZE; i += 0x1000) {
         copy_to_iram(IRAM_PAYLOAD_BASE + i, &g_reboot_payload[i], 0x1000);
     }
-    
     splSetConfig((SplConfigItem)65001, 2);
 }
 
-int reboot_to_payload(const char* path){
+int reboot_to_payload(const char* path, bool legacy){
     bool can_reboot = true;
     FILE *f;
     f = fopen(path, "rb");
     if (f == NULL) can_reboot = false;
     else {
-        printf("Can't open payload");
         fread(g_reboot_payload, 1, sizeof(g_reboot_payload), f);
         fclose(f);
     }
 
     if (can_reboot) {
-        printf("injecting payload");
-        inject_payload();
+        if (legacy) inject_payload_legacy();
+        else inject_payload();
     }
     if (can_reboot)
         splExit();

source/utils.cpp

@@ -210,16 +210,12 @@ std::vector<std::string> fetchPayloads(){
 }
 
 void shutDown(bool reboot){
-    //sync();
-    bpcInitialize();
-    if(reboot) bpcRebootSystem();
-    else bpcShutdownSystem();
-    bpcExit();
+    spsmInitialize();
+    spsmShutdown(reboot);
 }
 
 void rebootToPayload(const std::string& path) {
-    //sync();
-    reboot_to_payload(path.c_str());
+    reboot_to_payload(path.c_str(), CurrentCfw::running_cfw != CFW::ams);
 }
 
 std::string getLatestTag(const std::string& url){
@@ -247,26 +243,10 @@ std::string readVersion(const std::string& path){
 }
 
 bool isErista() {
-    u64 hwType;
-    Result rc = splGetConfig(SplConfigItem_HardwareType, &hwType);
-
-    if(R_FAILED(rc))
-        return true;
-
-    switch (hwType)
-    {
-        case 0:
-        case 1:
-            return true;
-        case 2:
-        case 3:
-        case 4:
-        case 5:
-            return false;
-        default:
-            return true;
-    }
-};
+    SetSysProductModel model;
+    setsysGetProductModel(&model);
+    return (model == SetSysProductModel_Nx || model == SetSysProductModel_Copper);
+}
 
 void removeSysmodulesFlags(const std::string&  directory) {
     for (const auto & e : std::filesystem::recursive_directory_iterator(directory)) {