3 files changed, 44 insertions, 63 deletions

diff --git a/arch/x86_64/src/boot/boot.s b/arch/x86_64/src/boot/boot.s
index 139fd1a..7932045 100644
--- a/arch/x86_64/src/boot/boot.s
+++ b/arch/x86_64/src/boot/boot.s
@@ -352,36 +352,6 @@ prepare_page_maps:
 .section .boot_text, "ax", @progbits
 .code64
 
-test_function:
-    cli
-    ret
-
-.global context_switch
-context_switch:
-    // ring 3 data with bottom 2 bits set for ring 3
-	mov $((4 * 16) | 3), %rax
-    mov %rax, %ds
-    mov %rax, %es
-    mov %rax, %fs
-    mov %rax, %gs
-    // SS is handled by iret https://wiki.osdev.org/Getting_to_Ring_3
-
-	// set up the stack frame iret expects
-	mov %rsp, %rax
-    // user data selector
-	push $((4 * 16) | 3)
-    // current rsp
-	push %rax
-    // push eflags
-	pushf
-    // push code selector (ring 3 code with bottom 2 bits set for ring 3)
-	push $((3 * 16) | 3)
-    // push instruction address to return to
-    lea [test_function], %rax
-	push %rax
-
-	iretq
-
 _transition_to_long_mode:
     xor %rax, %rax
     mov %rax, %ss
diff --git a/arch/x86_64/src/kernel/cpu/segment_register.cpp b/arch/x86_64/src/kernel/cpu/segment_register.cpp
index f70c558..d7857dd 100644
--- a/arch/x86_64/src/kernel/cpu/segment_register.cpp
+++ b/arch/x86_64/src/kernel/cpu/segment_register.cpp
@@ -21,7 +21,6 @@ namespace teachos::arch::kernel::cpu
   {
     asm volatile("xor %%rax, %%rax\n"
                  "mov %[input], %%ax\n"
-                 "mov %%rax, %%ss\n"
                  "mov %%rax, %%ds\n"
                  "mov %%rax, %%es\n"
                  "mov %%rax, %%fs\n"
diff --git a/arch/x86_64/src/kernel/main.cpp b/arch/x86_64/src/kernel/main.cpp
index daaf216..ac2591e 100644
--- a/arch/x86_64/src/kernel/main.cpp
+++ b/arch/x86_64/src/kernel/main.cpp
@@ -50,33 +50,48 @@ namespace teachos::arch::kernel
     delete test9;
   }
 
-  [[gnu::naked]]
-  auto push_code_segment(context_switching::interrupt_descriptor_table::segment_selector segment_selector_a,
-                         context_switching::interrupt_descriptor_table::segment_selector segment_selector_b) -> void
+  auto return_function() -> void
   {
-    asm volatile("push %%rbp\n"
-                 "push %[input]"
-                 : /* No output from call */
-                 : [input] "m"(segment_selector_a));
-    asm volatile("mov %[input], %%ax\n"
-                 "mov %%ax, %%ss\n"
-                 "mov %%ax, %%ds\n"
-                 "mov %%ax, %%es\n"
-                 "mov %%ax, %%fs\n"
-                 "mov %%ax, %%gs"
-                 : /* No output from call */
-                 : [input] "m"(segment_selector_b));
-    asm volatile("iretq"
-                 : /* No output from call */
-                 : /* No input to call */);
+    video::vga::text::write("User Mode!!!", video::vga::text::common_attributes::green_on_black);
   }
 
+  /**
+   * @brief Switch context into the mode defined in the segment selectors.
+   *
+   * Setup the stack IRETQ expects to switch the mode:
+   * 1. push data selector
+   * 2. push current stack pointer
+   * 3. push eflags
+   * 4. push code segment selector
+   * 5. push return address
+   *
+   * @param data_segment
+   * @param code_segment
+   * @param address
+   */
   [[gnu::naked]]
-  auto iret() -> void
+  auto switch_context(context_switching::interrupt_descriptor_table::segment_selector data_segment,
+                      context_switching::interrupt_descriptor_table::segment_selector code_segment, uint64_t address)
+      -> void
   {
-    asm volatile("iretq"
-                 : /* No output from call */
-                 : /* No input to call */);
+    asm volatile("mov %[data_segment], %%rax\n"
+                 "mov %%rax, %%ds\n"
+                 "mov %%rax, %%es\n"
+                 "mov %%rax, %%fs\n"
+                 "mov %%rax, %%gs\n"
+                 "mov %%rsp, %%rax\n"
+
+                 "push %[data_segment]\n"
+                 "push %%rax\n"
+                 "pushfq\n"
+                 "push %[code_segment]\n"
+                 "mov %[return_function], %%rax\n"
+                 "push %%rax\n"
+
+                 "iretq\n"
+                 :
+                 : [data_segment] "m"(data_segment), [code_segment] "m"(code_segment), [return_function] "r"(address)
+                 : "rax");
   }
 
   auto main() -> void
@@ -95,17 +110,14 @@ namespace teachos::arch::kernel
 
     decltype(auto) descriptor_tables = context_switching::initialize_descriptor_tables();
 
-    // - Clear NT flag in EFLAGS register (for far return)
-
-    // - Push return instruction pointer
-    // - Push return code segment selector
-    // context_switching::interrupt_descriptor_table::segment_selector user_code_segment_selector{
-    //     3U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
-    // context_switching::interrupt_descriptor_table::segment_selector user_data_segment_selector{
-    //     4U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
-    // push_code_segment(user_code_segment_selector, user_data_segment_selector);
+    context_switching::interrupt_descriptor_table::segment_selector user_code_segment_selector{
+        3U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
+    context_switching::interrupt_descriptor_table::segment_selector user_data_segment_selector{
+        4U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
 
-    boot::context_switch();
+    cpu::set_segment_registers(user_data_segment_selector);
+    switch_context(user_data_segment_selector, user_code_segment_selector,
+                   reinterpret_cast<uint64_t>(&return_function));
 
     (void)descriptor_tables;
   }