Merge pull request #102 from Freax13/enhancement/timer-irq

add preemption timer

common/constants/src/lib.rs

@@ -332,6 +332,10 @@ impl IntoIterator for PageRange<PhysFrame> {
     }
 }
 
+// Irq vectors used in the kernel.
+pub const TLB_VECTOR: u8 = 0x20;
+pub const TIMER_VECTOR: u8 = 0x21;
+
 #[cfg(test)]
 fn check_ranges<T>(ranges: &[PageRange<T>])
 where

common/tdx-types/src/tdcall.rs

@@ -73,6 +73,17 @@ impl MdFieldId {
         true,
     );
 
+    pub const TDVPS_TSC_DEADLINE: Self = Self::new(
+        89,
+        ElementSizeCode::SixtyFour,
+        0,
+        0,
+        false,
+        ContextCode::TdVcpu,
+        32,
+        true,
+    );
+
     #[allow(clippy::too_many_arguments)]
     const fn new(
         field_code: u32,
@@ -217,16 +228,6 @@ impl Apic {
         self.0[0x20 / 4].store(value, Ordering::SeqCst);
     }
 
-    /// Returns the highest priority requested interrupt.
-    pub fn pending_vectora_todo(&self) -> Option<u8> {
-        (0..8).rev().find_map(|i| {
-            let offset = 0x100 | (i * 16);
-            let idx = offset / 4;
-            let irr = self.0[idx].load(Ordering::SeqCst);
-            (irr != 0).then(|| i as u8 * 32 + 31 - irr.leading_zeros() as u8)
-        })
-    }
-
     /// Returns the highest priority requested interrupt.
     pub fn pending_vector(&self) -> Option<u8> {
         (0..8).rev().find_map(|i| {

common/tdx-types/src/vmexit.rs

@@ -2,3 +2,4 @@ pub const VMEXIT_REASON_CPUID_INSTRUCTION: u32 = 10;
 pub const VMEXIT_REASON_HLT_INSTRUCTION: u32 = 12;
 pub const VMEXIT_REASON_VMCALL_INSTRUCTION: u32 = 18;
 pub const VMEXIT_REASON_MSR_WRITE: u32 = 32;
+pub const VMEXIT_REASON_PREEMPTION_TIMER_EXPIRED: u32 = 52;

host/mushroom/Cargo.toml

@@ -24,7 +24,7 @@ constants = { workspace = true }
 loader = { workspace = true }
 log-types = { workspace = true, features = ["std"] }
 mushroom-verify = { workspace = true, optional = true }
-nix = { version = "0.29.0", features = ["fs", "ioctl", "mman", "pthread", "resource", "signal"] }
+nix = { version = "0.29.0", features = ["fs", "ioctl", "mman", "pthread", "resource", "signal", "time"] }
 profiler-types = { workspace = true }
 qgs-client = { workspace = true, optional = true }
 rand = "0.8.5"

host/mushroom/src/insecure.rs

@@ -16,13 +16,23 @@ use anyhow::{bail, Context, Result};
 use bit_field::BitField;
 use constants::{
     physical_address::{kernel, supervisor, DYNAMIC_2MIB},
-    INSECURE_SUPERVISOR_CALL_PORT, MAX_APS_COUNT,
+    INSECURE_SUPERVISOR_CALL_PORT, MAX_APS_COUNT, TIMER_VECTOR,
 };
 use loader::Input;
-use nix::sys::pthread::pthread_kill;
+use nix::{
+    sys::{
+        pthread::pthread_kill,
+        signal::SigEvent,
+        time::TimeSpec,
+        timer::{Timer, TimerSetTimeFlags},
+    },
+    time::ClockId,
+    unistd::gettid,
+};
 use snp_types::PageType;
 use supervisor_services::{SlotIndex, SupervisorCallNr};
 use tracing::info;
+use volatile::map_field;
 use x86_64::registers::{
     control::{Cr0Flags, Cr4Flags},
     model_specific::EferFlags,
@@ -45,6 +55,8 @@ static YMM_OFFSET: LazyLock<usize> = LazyLock::new(|| {
     res.ebx as usize
 });
 
+const TIMER_PERIOD: Duration = Duration::from_millis(10);
+
 /// Create the VM, load the kernel, init & input and run the APs.
 pub fn main(
     kvm_handle: &KvmHandle,
@@ -293,9 +305,40 @@ fn run_kernel_vcpu(
     let xsave_size = *KVM_XSAVE_SIZE.get().unwrap();
 
     let run_state = &run_states[usize::from(id)];
-    run_state.wait();
+    run_state.wait(Duration::MAX);
+
+    // Setup a timer to reguluarly kick the thread out of KVM_RUN.
+    let mut timer = Timer::new(
+        ClockId::CLOCK_MONOTONIC,
+        SigEvent::new(nix::sys::signal::SigevNotify::SigevThreadId {
+            signal: SIG_KICK,
+            thread_id: gettid().as_raw(),
+            si_value: 0,
+        }),
+    )?;
+    timer.set(
+        nix::sys::timer::Expiration::Interval(TimeSpec::from_duration(TIMER_PERIOD)),
+        TimerSetTimeFlags::empty(),
+    )?;
+    let mut last_timer_injection = Instant::now();
+    let mut in_service_timer_irq = false;
 
     while !run_state.is_stopped() {
+        // Check if we need to inject a timer interrupt.
+        if !in_service_timer_irq && last_timer_injection.elapsed() >= TIMER_PERIOD {
+            if map_field!(kvm_run.ready_for_interrupt_injection).read() != 0 {
+                map_field!(kvm_run.request_interrupt_window).write(0);
+
+                ap.interrupt(TIMER_VECTOR)?;
+
+                last_timer_injection = Instant::now();
+                in_service_timer_irq = true;
+            } else {
+                // Ask to be notified when the guest can receive an interrupt.
+                map_field!(kvm_run.request_interrupt_window).write(1);
+            }
+        }
+
         // Run the AP.
         let res = ap.run();
         match res {
@@ -321,7 +364,10 @@ fn run_kernel_vcpu(
                             run_state.kick();
                         }
                     }
-                    nr if nr == SupervisorCallNr::Halt as u64 => run_state.wait(),
+                    nr if nr == SupervisorCallNr::Halt as u64 => {
+                        let timeout = TIMER_PERIOD.saturating_sub(last_timer_injection.elapsed());
+                        run_state.wait(timeout);
+                    }
                     nr if nr == SupervisorCallNr::Kick as u64 => {
                         let index = regs.rdi as usize;
                         run_states[index].kick();
@@ -390,6 +436,12 @@ fn run_kernel_vcpu(
                     k
                 });
             }
+            KvmExit::WrMsr(msr) => match msr.index {
+                // EOI.
+                0x80b => in_service_timer_irq = false,
+                index => unimplemented!("unimplemented MSR write to {index:#x}"),
+            },
+            KvmExit::IrqWindowOpen => {}
             KvmExit::Interrupted => {}
             exit => {
                 let regs = ap.get_regs()?;
@@ -428,24 +480,28 @@ impl RunState {
         *self.running.lock().unwrap() == NextRunStateValue::Stopped
     }
 
-    pub fn wait(&self) {
+    pub fn wait(&self, timeout: Duration) {
         drop(
             self.condvar
-                .wait_while(self.running.lock().unwrap(), |state| match *state {
-                    NextRunStateValue::Halted => {
-                        // Keep waiting.
-                        true
-                    }
-                    NextRunStateValue::Ready => {
-                        // Consume the ready state and return.
-                        *state = NextRunStateValue::Halted;
-                        false
-                    }
-                    NextRunStateValue::Stopped => {
-                        // Don't update the state, but return.
-                        false
-                    }
-                })
+                .wait_timeout_while(
+                    self.running.lock().unwrap(),
+                    timeout,
+                    |state| match *state {
+                        NextRunStateValue::Halted => {
+                            // Keep waiting.
+                            true
+                        }
+                        NextRunStateValue::Ready => {
+                            // Consume the ready state and return.
+                            *state = NextRunStateValue::Halted;
+                            false
+                        }
+                        NextRunStateValue::Stopped => {
+                            // Don't update the state, but return.
+                            false
+                        }
+                    },
+                )
                 .unwrap(),
         );
     }

host/mushroom/src/kvm.rs

@@ -792,6 +792,17 @@ impl VcpuHandle {
         Ok(())
     }
 
+    pub fn interrupt(&self, vector: u8) -> Result<()> {
+        #[repr(transparent)]
+        struct KvmInterrupt(u32);
+        let kvm_interrupt_val = KvmInterrupt(u32::from(vector));
+
+        ioctl_write_ptr!(kvm_interrupt, KVMIO, 0x86, KvmInterrupt);
+        let res = unsafe { kvm_interrupt(self.fd.as_raw_fd(), &kvm_interrupt_val) };
+        res.context("failed to interrupt")?;
+        Ok(())
+    }
+
     #[cfg(feature = "tdx")]
     unsafe fn memory_encrypt_op_tdx<'a>(
         &self,

tee/kernel/src/exception.rs

@@ -7,11 +7,14 @@ use core::{
     ptr::null_mut,
 };
 
-use crate::memory::pagetable::flush::{tlb_shootdown_handler, TLB_VECTOR};
+use crate::memory::pagetable::flush::tlb_shootdown_handler;
 use crate::spin::lazy::Lazy;
-use crate::user::process::syscall::cpu_state::exception_entry;
+use crate::time;
+use crate::user::process::syscall::cpu_state::{exception_entry, interrupt_entry};
 use alloc::alloc::alloc;
+use constants::{TIMER_VECTOR, TLB_VECTOR};
 use log::{debug, error, trace};
+use x86_64::registers::model_specific::Msr;
 use x86_64::structures::gdt::SegmentSelector;
 use x86_64::{
     instructions::tables::load_tss,
@@ -146,6 +149,7 @@ pub fn load_idt() {
             .set_handler_fn(general_protection_fault_handler);
         idt.page_fault.set_handler_fn(page_fault_handler);
         idt[TLB_VECTOR].set_handler_fn(tlb_shootdown_handler);
+        idt[TIMER_VECTOR].set_handler_fn(timer_handler);
 
         idt[0x80]
             .set_handler_fn(int0x80_handler)
@@ -328,6 +332,35 @@ extern "x86-interrupt" fn double_fault_handler(frame: InterruptStackFrame, code:
     panic!("double fault {frame:x?} {code:x?}");
 }
 
+#[naked]
+extern "x86-interrupt" fn timer_handler(frame: InterruptStackFrame) {
+    unsafe {
+        naked_asm!(
+            "cld",
+            // Check whether the irq happened in userspace.
+            "test word ptr [rsp+8], 3",
+            "je {kernel_timer_handler}",
+
+            // Userspace code path:
+            "swapgs",
+            // Store the error code.
+            "mov byte ptr gs:[{VECTOR_OFFSET}], {TIMER_VECTOR}",
+            // Jump to the userspace exit point.
+            "jmp {interrupt_entry}",
+
+            kernel_timer_handler = sym kernel_timer_handler,
+            VECTOR_OFFSET = const offset_of!(PerCpu, vector),
+            TIMER_VECTOR = const TIMER_VECTOR,
+            interrupt_entry = sym interrupt_entry,
+        );
+    }
+}
+
+extern "x86-interrupt" fn kernel_timer_handler(_: InterruptStackFrame) {
+    time::try_fire_clocks();
+    eoi();
+}
+
 #[naked]
 extern "x86-interrupt" fn int0x80_handler(frame: InterruptStackFrame) {
     // The code that entered userspace stored addresses where execution should
@@ -343,3 +376,10 @@ extern "x86-interrupt" fn int0x80_handler(frame: InterruptStackFrame) {
         );
     }
 }
+
+/// Signal EOI.
+pub fn eoi() {
+    unsafe {
+        Msr::new(0x80b).write(0);
+    }
+}

tee/kernel/src/memory/pagetable/flush.rs

@@ -4,7 +4,7 @@ use core::{
 };
 
 use bit_field::{BitArray, BitField};
-use constants::{ApBitmap, ApIndex, AtomicApBitmap, MAX_APS_COUNT};
+use constants::{ApBitmap, ApIndex, AtomicApBitmap, MAX_APS_COUNT, TLB_VECTOR};
 use x86_64::{
     instructions::tlb::{self, InvPicdCommand, Invlpgb},
     registers::{
@@ -14,12 +14,10 @@ use x86_64::{
     structures::{idt::InterruptStackFrame, paging::Page},
 };
 
-use crate::{per_cpu::PerCpu, spin::lazy::Lazy};
+use crate::{exception::eoi, per_cpu::PerCpu, spin::lazy::Lazy};
 
 use super::ActivePageTableGuard;
 
-pub const TLB_VECTOR: u8 = 0x20;
-
 static INVLPGB: Lazy<Option<Invlpgb>> = Lazy::new(Invlpgb::new);
 
 static ACTIVE_APS: AtomicApBitmap = AtomicApBitmap::empty();
@@ -80,10 +78,7 @@ pub extern "x86-interrupt" fn tlb_shootdown_handler(_: InterruptStackFrame) {
 
     process_flushes(idx);
 
-    // Signal EOI.
-    unsafe {
-        Msr::new(0x80b).write(0);
-    }
+    eoi();
 }
 
 fn send_tlb_ipis(aps: ApBitmap) {

tee/kernel/src/rt.rs

@@ -3,7 +3,7 @@ use core::{
     future::Future,
     panic::Location,
     pin::Pin,
-    task::{Context, Waker},
+    task::{Context, Poll, Waker},
 };
 
 use crate::{spin::mutex::Mutex, user::schedule_vcpu};
@@ -140,3 +140,25 @@ enum TaskState {
     /// The task has finished.
     Done,
 }
+
+pub async fn r#yield() {
+    struct Yield {
+        polled: bool,
+    }
+
+    impl Future for Yield {
+        type Output = ();
+
+        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            if self.polled {
+                Poll::Ready(())
+            } else {
+                self.polled = true;
+                cx.waker().wake_by_ref();
+                Poll::Pending
+            }
+        }
+    }
+
+    Yield { polled: false }.await
+}

tee/kernel/src/time.rs

@@ -119,6 +119,15 @@ pub fn now(clock: ClockId) -> Timespec {
     STATE.lock().read_clock(clock)
 }
 
+pub fn try_fire_clocks() {
+    let Some(mut guard) = STATE.try_lock() else {
+        // Some other thread is already using the state. Don't do anything
+        // now.
+        return;
+    };
+    guard.fire_clocks();
+}
+
 pub fn set(clock: ClockId, time: Timespec) -> Result<()> {
     match clock {
         ClockId::Realtime => STATE.lock().set_real_time(time),