Skip to main content

strat9_kernel/arch/x86_64/
tss.rs

1//! Task State Segment (TSS) for Strat9-OS
2//!
3//! The TSS is required for:
4//! - Interrupt Stack Table (IST) entries for safe exception handling
5//! - Ring 3 -> Ring 0 stack switching (privilege_stack_table[0] = rsp0)
6
7use core::{
8    mem::MaybeUninit,
9    panic::Location,
10    sync::atomic::{AtomicBool, Ordering},
11};
12use x86_64::{structures::tss::TaskStateSegment, VirtAddr};
13
14#[repr(C, packed)]
15struct DescriptorTableRegister {
16    limit: u16,
17    base: u64,
18}
19
20#[derive(Clone, Copy, Debug)]
21pub struct LoadedTssInfo {
22    pub tr_selector: u16,
23    pub tss_base: u64,
24    pub rsp0: u64,
25}
26
27/// IST index used for the double fault handler
28pub const DOUBLE_FAULT_IST_INDEX: u16 = 0;
29
30/// IST stack size (20 KB)
31const IST_STACK_SIZE: usize = 4096 * 5;
32
33/// Static IST stacks for double fault handler (per-CPU)
34static mut IST_STACKS: [[u8; IST_STACK_SIZE]; crate::arch::x86_64::percpu::MAX_CPUS] =
35    [[0; IST_STACK_SIZE]; crate::arch::x86_64::percpu::MAX_CPUS];
36
37/// Per-CPU TSS storage
38static mut TSS: [MaybeUninit<TaskStateSegment>; crate::arch::x86_64::percpu::MAX_CPUS] =
39    [const { MaybeUninit::uninit() }; crate::arch::x86_64::percpu::MAX_CPUS];
40
41static TSS_INIT: [AtomicBool; crate::arch::x86_64::percpu::MAX_CPUS] =
42    [const { AtomicBool::new(false) }; crate::arch::x86_64::percpu::MAX_CPUS];
43
44/// Initialize the TSS with IST entries
45///
46/// Must be called before `gdt::init()` since the GDT references the TSS.
47pub fn init() {
48    init_cpu(0);
49}
50
51/// Initialize the TSS for a given CPU index.
52pub fn init_cpu(cpu_index: usize) {
53    // Bounds check: prevent OOB access into static arrays before any unsafe.
54    assert!(
55        cpu_index < crate::arch::x86_64::percpu::MAX_CPUS,
56        "TSS init_cpu: cpu_index {} >= MAX_CPUS {}",
57        cpu_index,
58        crate::arch::x86_64::percpu::MAX_CPUS,
59    );
60    // SAFETY: Called during init (BSP) or AP bring-up before interrupts are enabled on that CPU.
61    unsafe {
62        let stack_ptr = &raw const IST_STACKS[cpu_index] as *const u8;
63        let stack_end = VirtAddr::from_ptr(stack_ptr) + IST_STACK_SIZE as u64;
64        let mut tss = TaskStateSegment::new();
65        tss.interrupt_stack_table[DOUBLE_FAULT_IST_INDEX as usize] = stack_end;
66
67        TSS[cpu_index].write(tss);
68        TSS_INIT[cpu_index].store(true, Ordering::Release);
69
70        let ist_addr = VirtAddr::from_ptr(stack_ptr);
71        // Raw e9 output : format_args! + Port write hangs before IDT is loaded.
72        for &b in b"TSS init OK\n" {
73            unsafe {
74                core::arch::asm!("out 0xe9, al", in("al") b, options(nomem, nostack));
75            }
76        }
77    }
78}
79
80/// Get a reference to the TSS for a given CPU index (for GDT descriptor creation).
81pub fn tss_for(cpu_index: usize) -> &'static TaskStateSegment {
82    assert!(
83        cpu_index < crate::arch::x86_64::percpu::MAX_CPUS,
84        "tss_for: cpu_index {} >= MAX_CPUS",
85        cpu_index,
86    );
87    if !TSS_INIT[cpu_index].load(Ordering::Acquire) {
88        panic!("TSS for CPU{} not initialized", cpu_index);
89    }
90    // SAFETY: TSS was initialized in init_cpu and lives for 'static.
91    unsafe { &*TSS[cpu_index].as_ptr() }
92}
93
94/// Return TSS.rsp0 for a specific CPU, if its TSS is initialized.
95pub fn kernel_stack_for(cpu_index: usize) -> Option<VirtAddr> {
96    if cpu_index >= crate::arch::x86_64::percpu::MAX_CPUS {
97        return None;
98    }
99    if !TSS_INIT[cpu_index].load(Ordering::Acquire) {
100        return None;
101    }
102    // SAFETY: The TSS for this CPU is initialized and lives for the whole kernel lifetime.
103    unsafe {
104        let tss = &*TSS[cpu_index].as_ptr();
105        Some(tss.privilege_stack_table[0])
106    }
107}
108
109/// Read the TSS currently loaded in TR by decoding the live GDT entry.
110pub fn loaded_tss_info() -> Option<LoadedTssInfo> {
111    let mut gdtr = DescriptorTableRegister { limit: 0, base: 0 };
112    let tr_selector: u16;
113    // SAFETY: `sgdt`/`str` are privileged register reads with no side effect.
114    unsafe {
115        core::arch::asm!(
116            "sgdt [{}]",
117            in(reg) &mut gdtr,
118            options(nostack, preserves_flags),
119        );
120        core::arch::asm!(
121            "str {0:x}",
122            out(reg) tr_selector,
123            options(nostack, nomem, preserves_flags),
124        );
125    }
126
127    if tr_selector == 0 {
128        return None;
129    }
130
131    let entry_offset = (tr_selector & !0x7) as usize;
132    if entry_offset + 16 > gdtr.limit as usize + 1 {
133        return None;
134    }
135
136    // SAFETY: The GDTR base/limit were read from the CPU and bounds-checked above.
137    let (low, high) = unsafe {
138        let entry_ptr = (gdtr.base + entry_offset as u64) as *const u64;
139        (
140            core::ptr::read_unaligned(entry_ptr),
141            core::ptr::read_unaligned(entry_ptr.add(1)),
142        )
143    };
144
145    let base_low =
146        ((low >> 16) & 0xFFFF) | (((low >> 32) & 0xFF) << 16) | (((low >> 56) & 0xFF) << 24);
147    let tss_base = base_low | (high << 32);
148    if tss_base == 0 {
149        return None;
150    }
151
152    // SAFETY: The live TSS base comes from the CPU's TSS descriptor.
153    let rsp0 = unsafe {
154        let tss = &*(tss_base as *const TaskStateSegment);
155        tss.privilege_stack_table[0].as_u64()
156    };
157
158    Some(LoadedTssInfo {
159        tr_selector,
160        tss_base,
161        rsp0,
162    })
163}
164
165/// Update TSS.rsp0 : the kernel stack pointer used when transitioning
166/// from Ring 3 to Ring 0 on interrupt/syscall.
167///
168/// Called on every context switch to point to the new task's kernel stack top.
169#[track_caller]
170pub fn set_kernel_stack(stack_top: VirtAddr) {
171    let cpu_index = crate::arch::x86_64::percpu::current_cpu_index();
172    set_kernel_stack_for(cpu_index, stack_top);
173}
174
175/// Update TSS.rsp0 for a specific CPU index.
176#[track_caller]
177pub fn set_kernel_stack_for(cpu_index: usize, stack_top: VirtAddr) {
178    if cpu_index >= crate::arch::x86_64::percpu::MAX_CPUS {
179        log::warn!("set_kernel_stack_for: cpu_index {} out of range", cpu_index);
180        return;
181    }
182    // SAFETY: privilege_stack_table[0] is a VirtAddr (u64), writes are atomic on x86_64.
183    // Called with interrupts disabled or from the scheduler with lock held.
184    if !TSS_INIT[cpu_index].load(Ordering::Acquire) {
185        return;
186    }
187    let caller = Location::caller();
188    unsafe {
189        let tss = &raw mut *TSS[cpu_index].as_mut_ptr();
190        let old_stack_top = (*tss).privilege_stack_table[0];
191        (*tss).privilege_stack_table[0] = stack_top;
192        crate::e9_println!(
193            "[tss-set] cpu={} old={:#x} new={:#x} caller={}:{}",
194            cpu_index,
195            old_stack_top.as_u64(),
196            stack_top.as_u64(),
197            caller.file(),
198            caller.line()
199        );
200    }
201}