From 04e0f8776c26930d7dc8015e53914b11bf1929c1 Mon Sep 17 00:00:00 2001
From: Fernando Sahmkow <fsahmkow27@gmail.com>
Date: Thu, 27 Feb 2020 19:12:41 -0400
Subject: [PATCH] General: Add better safety for JIT use.

---
 src/core/arm/arm_interface.h      | 10 ++++++++++
 src/core/cpu_manager.cpp          | 25 +++++++++++++++++++------
 src/core/cpu_manager.h            |  2 ++
 src/core/hle/kernel/scheduler.cpp |  3 +++
 src/core/hle/kernel/svc.cpp       |  6 +++++-
 5 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/core/arm/arm_interface.h b/src/core/arm/arm_interface.h
index 87a1c29cc..be9f3703a 100644
--- a/src/core/arm/arm_interface.h
+++ b/src/core/arm/arm_interface.h
@@ -6,6 +6,7 @@
 
 #include <array>
 #include <vector>
+#include <mutex>
 #include "common/common_types.h"
 
 namespace Common {
@@ -164,6 +165,14 @@ public:
         std::string name;
     };
 
+    void Lock() {
+        guard.lock();
+    }
+
+    void Unlock() {
+        guard.unlock();
+    }
+
     std::vector<BacktraceEntry> GetBacktrace() const;
 
     /// fp (= r29) points to the last frame record.
@@ -178,6 +187,7 @@ protected:
     /// System context that this ARM interface is running under.
     System& system;
     CPUInterruptHandler& interrupt_handler;
+    std::mutex guard;
 };
 
 } // namespace Core
diff --git a/src/core/cpu_manager.cpp b/src/core/cpu_manager.cpp
index 904aacd97..9a261968a 100644
--- a/src/core/cpu_manager.cpp
+++ b/src/core/cpu_manager.cpp
@@ -46,6 +46,11 @@ void CpuManager::GuestThreadFunction(void* cpu_manager_) {
     cpu_manager->RunGuestThread();
 }
 
+void CpuManager::GuestRewindFunction(void* cpu_manager_) {
+    CpuManager* cpu_manager = static_cast<CpuManager*>(cpu_manager_);
+    cpu_manager->RunGuestLoop();
+}
+
 void CpuManager::IdleThreadFunction(void* cpu_manager_) {
     CpuManager* cpu_manager = static_cast<CpuManager*>(cpu_manager_);
     cpu_manager->RunIdleThread();
@@ -78,14 +83,22 @@ void CpuManager::RunGuestThread() {
         auto& sched = kernel.CurrentScheduler();
         sched.OnThreadStart();
     }
+    RunGuestLoop();
+}
+
+void CpuManager::RunGuestLoop() {
+    auto& kernel = system.Kernel();
+    auto* thread = kernel.CurrentScheduler().GetCurrentThread();
+    auto host_context = thread->GetHostContext();
+    host_context->SetRewindPoint(std::function<void(void*)>(GuestRewindFunction), this);
+    host_context.reset();
     while (true) {
-        auto* physical_core = &kernel.CurrentPhysicalCore();
-        while (!physical_core->IsInterrupted()) {
-            physical_core->Run();
-            physical_core = &kernel.CurrentPhysicalCore();
+        auto& physical_core = kernel.CurrentPhysicalCore();
+        while (!physical_core.IsInterrupted()) {
+            physical_core.Run();
         }
-        physical_core->ClearExclusive();
-        auto& scheduler = physical_core->Scheduler();
+        physical_core.ClearExclusive();
+        auto& scheduler = physical_core.Scheduler();
         scheduler.TryDoContextSwitch();
     }
 }
diff --git a/src/core/cpu_manager.h b/src/core/cpu_manager.h
index 8103ae857..e83ab20f9 100644
--- a/src/core/cpu_manager.h
+++ b/src/core/cpu_manager.h
@@ -42,10 +42,12 @@ public:
 
 private:
     static void GuestThreadFunction(void* cpu_manager);
+    static void GuestRewindFunction(void* cpu_manager);
     static void IdleThreadFunction(void* cpu_manager);
     static void SuspendThreadFunction(void* cpu_manager);
 
     void RunGuestThread();
+    void RunGuestLoop();
     void RunIdleThread();
     void RunSuspendThread();
 
diff --git a/src/core/hle/kernel/scheduler.cpp b/src/core/hle/kernel/scheduler.cpp
index 758fa8188..727d2e6cc 100644
--- a/src/core/hle/kernel/scheduler.cpp
+++ b/src/core/hle/kernel/scheduler.cpp
@@ -581,6 +581,7 @@ void Scheduler::SwitchContextStep2() {
 
     if (new_thread) {
         new_thread->context_guard.lock();
+        cpu_core.Lock();
         ASSERT_MSG(new_thread->GetProcessorID() == s32(this->core_id),
                    "Thread must be assigned to this core.");
         ASSERT_MSG(new_thread->GetStatus() == ThreadStatus::Ready,
@@ -601,6 +602,7 @@ void Scheduler::SwitchContextStep2() {
             cpu_core.LoadContext(new_thread->GetContext64());
             cpu_core.SetTlsAddress(new_thread->GetTLSAddress());
             cpu_core.SetTPIDR_EL0(new_thread->GetTPIDR_EL0());
+            cpu_core.ClearExclusiveState();
         }
     } else {
         current_thread = nullptr;
@@ -639,6 +641,7 @@ void Scheduler::SwitchContext() {
         }
         previous_thread->SetIsRunning(false);
         previous_thread->context_guard.unlock();
+        cpu_core.Unlock();
     }
 
     std::shared_ptr<Common::Fiber> old_context;
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 1e6c60d78..b535593c7 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -10,6 +10,7 @@
 
 #include "common/alignment.h"
 #include "common/assert.h"
+#include "common/fiber.h"
 #include "common/logging/log.h"
 #include "common/microprofile.h"
 #include "common/string_util.h"
@@ -2468,7 +2469,10 @@ void Call(Core::System& system, u32 immediate) {
     }
     auto& physical_core_2 = system.CurrentPhysicalCore();
     if (physical_core.CoreIndex() != physical_core_2.CoreIndex()) {
-        physical_core.Stop();
+        LOG_CRITICAL(Kernel_SVC, "Rewinding");
+        auto* thread = physical_core_2.Scheduler().GetCurrentThread();
+        auto* host_context = thread->GetHostContext().get();
+        host_context->Rewind();
     }
 }