From 9f44fe6c4e62b128b0e65acc6a406fecabaad81f Mon Sep 17 00:00:00 2001
From: mintsuki <mintsuki@protonmail.com>
Date: Fri, 18 Mar 2022 01:47:04 +0100
Subject: [PATCH] smp: Make it higher half aware

---
 common/protos/limine.c         |  3 ++-
 common/protos/stivale2.c       |  3 ++-
 common/sys/smp.c               | 11 +++++++----
 common/sys/smp.h               |  3 ++-
 common/sys/smp_trampoline.real | 10 +++++++++-
 5 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/common/protos/limine.c b/common/protos/limine.c
index bbbc730b..2c38e6aa 100644
--- a/common/protos/limine.c
+++ b/common/protos/limine.c
@@ -441,7 +441,8 @@ FEAT_START
     smp_info = init_smp(0, (void **)&smp_array,
                         &cpu_count, &bsp_lapic_id,
                         true, want_5lv,
-                        pagemap, smp_request->flags & LIMINE_SMP_X2APIC, true);
+                        pagemap, smp_request->flags & LIMINE_SMP_X2APIC, true,
+                        direct_map_offset);
 
     if (smp_info == NULL) {
         break;
diff --git a/common/protos/stivale2.c b/common/protos/stivale2.c
index 68d728fb..95ab5e99 100644
--- a/common/protos/stivale2.c
+++ b/common/protos/stivale2.c
@@ -763,7 +763,8 @@ have_tm_tag:;
         smp_info = init_smp(sizeof(struct stivale2_struct_tag_smp), (void **)&tag,
                             &cpu_count, &bsp_lapic_id,
                             bits == 64, want_5lv,
-                            pagemap, smp_hdr_tag->flags & 1, want_pmrs);
+                            pagemap, smp_hdr_tag->flags & 1, want_pmrs,
+                            stivale2_hdr.flags & (1 << 1) ? direct_map_offset : 0);
 
         if (smp_info != NULL) {
             tag->tag.identifier = STIVALE2_STRUCT_TAG_SMP_ID;
diff --git a/common/sys/smp.c b/common/sys/smp.c
index 963adf2d..3c60518c 100644
--- a/common/sys/smp.c
+++ b/common/sys/smp.c
@@ -48,12 +48,13 @@ struct trampoline_passed_info {
     uint32_t smp_tpl_pagemap;
     uint32_t smp_tpl_info_struct;
     struct gdtr smp_tpl_gdt;
+    uint64_t smp_tpl_hhdm;
 } __attribute__((packed));
 
 static bool smp_start_ap(uint32_t lapic_id, struct gdtr *gdtr,
                          struct smp_information *info_struct,
                          bool longmode, bool lv5, uint32_t pagemap,
-                         bool x2apic, bool nx) {
+                         bool x2apic, bool nx, uint64_t hhdm) {
     size_t trampoline_size = (size_t)_binary_smp_trampoline_bin_end
                            - (size_t)_binary_smp_trampoline_bin_start;
 
@@ -80,6 +81,7 @@ static bool smp_start_ap(uint32_t lapic_id, struct gdtr *gdtr,
                         | (uint32_t)longmode;
     passed_info->smp_tpl_gdt         = *gdtr;
     passed_info->smp_tpl_booted_flag = 0;
+    passed_info->smp_tpl_hhdm = hhdm;
 
     asm volatile ("" ::: "memory");
 
@@ -119,7 +121,8 @@ struct smp_information *init_smp(size_t    header_hack_size,
                                  bool      lv5,
                                  pagemap_t pagemap,
                                  bool      x2apic,
-                                 bool      nx) {
+                                 bool      nx,
+                                 uint64_t  hhdm) {
     if (!lapic_check())
         return NULL;
 
@@ -224,7 +227,7 @@ struct smp_information *init_smp(size_t    header_hack_size,
                 // Try to start the AP
                 if (!smp_start_ap(lapic->lapic_id, &gdtr, info_struct,
                                   longmode, lv5, (uintptr_t)pagemap.top_level,
-                                  x2apic, nx)) {
+                                  x2apic, nx, hhdm)) {
                     print("smp: FAILED to bring-up AP\n");
                     continue;
                 }
@@ -261,7 +264,7 @@ struct smp_information *init_smp(size_t    header_hack_size,
                 // Try to start the AP
                 if (!smp_start_ap(x2lapic->x2apic_id, &gdtr, info_struct,
                                   longmode, lv5, (uintptr_t)pagemap.top_level,
-                                  true, nx)) {
+                                  true, nx, hhdm)) {
                     print("smp: FAILED to bring-up AP\n");
                     continue;
                 }
diff --git a/common/sys/smp.h b/common/sys/smp.h
index efc0f7c6..df272871 100644
--- a/common/sys/smp.h
+++ b/common/sys/smp.h
@@ -22,6 +22,7 @@ struct smp_information *init_smp(size_t    header_hack_size,
                                  bool      lv5,
                                  pagemap_t pagemap,
                                  bool      x2apic,
-                                 bool      nx);
+                                 bool      nx,
+                                 uint64_t  hhdm);
 
 #endif
diff --git a/common/sys/smp_trampoline.real b/common/sys/smp_trampoline.real
index fa8b173e..2d672e20 100644
--- a/common/sys/smp_trampoline.real
+++ b/common/sys/smp_trampoline.real
@@ -101,7 +101,13 @@ smp_trampoline:
     wrmsr
 
   .nonx:
-    jmp parking64
+    mov rax, qword [rbx + passed_info.hhdm]
+    add qword [rbx + passed_info.gdtr + 2], rax
+    lgdt [rbx + passed_info.gdtr]
+
+    lea rax, [rax + rbx + parking64]
+
+    jmp rax
 
 bits 32
 parking32:
@@ -183,3 +189,5 @@ passed_info:
     .gdtr:
         dw 0
         dq 0
+    .hhdm:
+        dq 0