1 files changed, 43 insertions, 5 deletions

diff --git a/arch/x86_64/src/kernel/cpu/segment_register.cpp b/arch/x86_64/src/kernel/cpu/segment_register.cpp
index cb367b6..44b4eff 100644
--- a/arch/x86_64/src/kernel/cpu/segment_register.cpp
+++ b/arch/x86_64/src/kernel/cpu/segment_register.cpp
@@ -6,7 +6,7 @@
 namespace teachos::arch::kernel::cpu
 {
   [[gnu::naked]]
-  auto reload_segment_registers() -> void
+  auto reload_data_segment_registers() -> void
   {
     asm volatile("xor %rax, %rax\n"
                  "mov %rax, %ss\n"
@@ -18,7 +18,8 @@ namespace teachos::arch::kernel::cpu
   }
 
   [[gnu::naked]]
-  auto set_segment_registers(context_switching::interrupt_descriptor_table::segment_selector segment_selector) -> void
+  auto set_data_segment_registers(context_switching::interrupt_descriptor_table::segment_selector segment_selector)
+      -> void
   {
     asm volatile("xor %%rax, %%rax\n"
                  "mov %[input], %%ax\n"
@@ -38,7 +39,8 @@ namespace teachos::arch::kernel::cpu
     return current_value;
   }
 
-  auto validate_data_segment_registers() -> context_switching::interrupt_descriptor_table::segment_selector
+  auto validate_data_segment_registers(context_switching::interrupt_descriptor_table::segment_selector data_segment)
+      -> void
   {
     context_switching::interrupt_descriptor_table::segment_selector ss;
     context_switching::interrupt_descriptor_table::segment_selector ds;
@@ -54,10 +56,46 @@ namespace teachos::arch::kernel::cpu
         "mov %%gs, %[gs_output]\n"
         : [ss_output] "=r"(ss), [ds_output] "=r"(ds), [es_output] "=r"(es), [fs_output] "=r"(fs), [gs_output] "=r"(gs));
 
-    auto result = ss == ds && ss == es && ss == fs && ss == gs;
+    auto result = (ss == ds && ss == es && ss == fs && ss == gs);
     exception_handling::assert(result, "[Segment Register] Values in data register are not the same.");
+    result = (ss == data_segment);
+    exception_handling::assert(
+        result, "[Segment Register] Expected Data Segment is not the same as the value in the Stack Segment register.");
+  }
+
+  auto validate_code_segment_register(context_switching::interrupt_descriptor_table::segment_selector code_segment)
+      -> void
+  {
+    auto const cs = read_code_segment_register();
+    exception_handling::assert(
+        cs == code_segment,
+        "[Segment Register] Expected Code Segment is not the same as the value in the Code Segment register.");
+  }
+
+  auto validate_segment_registers(context_switching::interrupt_descriptor_table::segment_selector data_segment,
+                                  context_switching::interrupt_descriptor_table::segment_selector code_segment) -> void
+  {
+    validate_data_segment_registers(data_segment);
+    validate_code_segment_register(code_segment);
+  }
+
+  [[gnu::naked]]
+  auto set_code_segment_register(context_switching::interrupt_descriptor_table::segment_selector data_segment,
+                                 context_switching::interrupt_descriptor_table::segment_selector code_segment,
+                                 uint64_t address) -> void
+  {
+    asm volatile("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"
 
-    return ss;
+                 "iretq\n"
+                 :
+                 : [data_segment] "m"(data_segment), [code_segment] "m"(code_segment), [return_function] "r"(address)
+                 : "rax");
   }
 
 }  // namespace teachos::arch::kernel::cpu