deps: upgrade td-shim

Signed-off-by: Jiaqi Gao <[email protected]>

src/attestation/src/ghci.rs

@@ -5,8 +5,8 @@
 use core::sync::atomic::{AtomicU8, Ordering};
 use core::{ffi::c_void, slice::from_raw_parts_mut};
 use td_payload::arch::apic::{disable, enable_and_hlt};
-use td_payload::arch::idt::register;
-use td_payload::{interrupt_handler_template, mm::shared::SharedMemory};
+use td_payload::arch::idt::{register_interrupt_callback, InterruptCallback, InterruptStack};
+use td_payload::mm::shared::SharedMemory;
 use tdx_tdcall::tdx::tdvmcall_get_quote;
 
 use crate::binding::AttestLibError;
@@ -46,13 +46,20 @@ pub extern "C" fn servtd_get_quote(tdquote_req_buf: *mut c_void, len: u64) -> i3
     0
 }
 
-interrupt_handler_template!(vmm_notification, _stack, {
+fn vmm_notification(_: &mut InterruptStack) {
     NOTIFIER.store(NOTIFY_VALUE, Ordering::SeqCst);
-});
+}
 
 pub fn set_vmm_notification() {
     // Setup interrupt handler
-    register(NOTIFY_VECTOR, vmm_notification);
+    if register_interrupt_callback(
+        NOTIFY_VECTOR as usize,
+        InterruptCallback::new(vmm_notification),
+    )
+    .is_err()
+    {
+        panic!("Fail to setup interrupt callback for VMM notify\n");
+    }
 
     // Setup event notifier
     if tdx_tdcall::tdx::tdvmcall_setup_event_notify(NOTIFY_VECTOR as u64).is_err() {

src/devices/virtio_serial/src/event.rs

@@ -2,15 +2,18 @@
 //
 // SPDX-License-Identifier: BSD-2-Clause-Patent
 
-use core::sync::atomic::{AtomicBool, Ordering};
+use core::{
+    result::Result,
+    sync::atomic::{AtomicBool, Ordering},
+};
 pub use td_payload::arch::apic::*;
-use td_payload::arch::idt::register;
-pub use td_payload::interrupt_handler_template;
+use td_payload::arch::idt::{register_interrupt_callback, InterruptCallback, InterruptStack};
 
-use crate::Timer;
+use crate::{Timer, VirtioSerialError};
 
-pub fn register_callback(vector: u8, cb: unsafe extern "C" fn()) {
-    register(vector, cb);
+pub fn register_callback(vector: u8, cb: fn(&mut InterruptStack)) -> Result<(), VirtioSerialError> {
+    register_interrupt_callback(vector as usize, InterruptCallback::new(cb))
+        .map_err(|_| VirtioSerialError::Interrupt)
 }
 
 pub fn wait_for_event(event_flag: &AtomicBool, timer: &dyn Timer) -> bool {

src/devices/virtio_serial/src/lib.rs

@@ -22,6 +22,7 @@ use core::slice::{from_raw_parts, from_raw_parts_mut};
 use core::sync::atomic::{AtomicBool, Ordering};
 use lazy_static::lazy_static;
 use spin::{Mutex, Once};
+use td_payload::arch::idt::InterruptStack;
 use virtio::{consts::*, virtqueue::*, VirtioError, VirtioTransport};
 
 use rust_std_stub::{error, io};
@@ -65,6 +66,8 @@ pub enum VirtioSerialError {
     PortNotAvailable(u32),
     // The port is already occupied
     PortAlreadyUsed(u32),
+    /// Configure device interrupt
+    Interrupt,
 }
 
 impl Display for VirtioSerialError {
@@ -78,6 +81,7 @@ impl Display for VirtioSerialError {
             VirtioSerialError::Timeout => write!(f, "Timeout"),
             VirtioSerialError::PortNotAvailable(e) => write!(f, "PortNotAvailable: 0x{:x}", e),
             VirtioSerialError::PortAlreadyUsed(e) => write!(f, "PortAlreadyUsed: 0x{:x}", e),
+            VirtioSerialError::Interrupt => write!(f, "Interrupt"),
         }
     }
 }
@@ -393,7 +397,7 @@ impl VirtioSerial {
     }
 
     fn init_notification(&mut self) -> Result<()> {
-        register_callback(IRQ_VECTOR, serial_event_callback);
+        register_callback(IRQ_VECTOR, serial_event_callback)?;
         let transport = self.virtio.as_mut();
 
         let irq_index = transport.set_interrupt_vector(IRQ_VECTOR)?;
@@ -873,6 +877,6 @@ pub(crate) fn align_up(size: usize) -> usize {
     (size & !(PAGE_SIZE - 1)) + if size % PAGE_SIZE != 0 { PAGE_SIZE } else { 0 }
 }
 
-interrupt_handler_template!(serial_event_callback, _stack, {
+fn serial_event_callback(_: &mut InterruptStack) {
     IRQ_FLAG.store(true, Ordering::SeqCst);
-});
+}

src/devices/vsock/src/transport/event.rs

@@ -2,19 +2,31 @@
 //
 // SPDX-License-Identifier: BSD-2-Clause-Patent
 
+use super::VsockTransportError;
+
+use core::result::Result;
 use core::sync::atomic::{AtomicBool, Ordering};
-pub use td_payload::arch::apic::*;
-use td_payload::arch::idt::register;
-pub use td_payload::interrupt_handler_template;
+use td_payload::arch::apic::*;
+use td_payload::arch::idt::{register_interrupt_callback, InterruptCallback, InterruptStack};
 
 use crate::VsockTimeout;
 
 #[cfg(not(feature = "fuzz"))]
-pub fn register_callback(vector: u8, cb: unsafe extern "C" fn()) {
-    register(vector, cb);
+pub fn register_callback(
+    vector: u8,
+    cb: fn(&mut InterruptStack),
+) -> Result<(), VsockTransportError> {
+    register_interrupt_callback(vector as usize, InterruptCallback::new(cb))
+        .map_err(|_| VsockTransportError::Interrupt)
 }
+
 #[cfg(feature = "fuzz")]
-pub fn register_callback(vector: u8, cb: unsafe extern "C" fn()) {}
+pub fn register_callback(
+    vector: u8,
+    cb: fn(&mut InterruptStack),
+) -> Result<(), VsockTransportError> {
+    Ok(())
+}
 
 pub fn wait_for_event(event_flag: &AtomicBool, timer: &dyn VsockTimeout) -> bool {
     while !event_flag.load(Ordering::SeqCst) {

src/devices/vsock/src/transport/mod.rs

@@ -35,6 +35,7 @@ pub enum VsockTransportError {
     Timeout,
     InvalidParameter,
     InvalidVsockPacket,
+    Interrupt,
 }
 
 impl Display for VsockTransportError {
@@ -50,6 +51,7 @@ impl Display for VsockTransportError {
             VsockTransportError::Virtio(e) => write!(f, "Virtio: {}", e),
             #[cfg(feature = "vmcall-vsock")]
             VsockTransportError::Vmcall(_) => write!(f, "Vmcall"),
+            VsockTransportError::Interrupt => write!(f, "Interrupt"),
         }
     }
 }

src/devices/vsock/src/transport/virtio_pci.rs

@@ -18,6 +18,7 @@ use alloc::vec::Vec;
 use core::cell::RefCell;
 use core::ptr::{slice_from_raw_parts, slice_from_raw_parts_mut};
 use core::sync::atomic::{AtomicBool, Ordering};
+use td_payload::arch::idt::InterruptStack;
 use virtio::virtqueue::{VirtQueue, VirtQueueLayout, VirtqueueBuf};
 use virtio::{consts::*, VirtioError, VirtioTransport};
 
@@ -129,9 +130,9 @@ impl VirtioVsock {
         let irq_index_tx = transport.set_interrupt_vector(TX_VECTOR)?;
         Self::set_queue_notify(transport, QUEUE_TX, irq_index_tx)?;
 
-        register_callback(RX_VECTOR, rx_callback);
-        register_callback(TX_VECTOR, tx_callback);
-        register_callback(CONFIG_VECTOR, config_callback);
+        register_callback(RX_VECTOR, rx_callback)?;
+        register_callback(TX_VECTOR, tx_callback)?;
+        register_callback(CONFIG_VECTOR, config_callback)?;
 
         Ok(Self {
             virtio_transport,
@@ -444,14 +445,14 @@ impl Drop for VirtioVsock {
     }
 }
 
-interrupt_handler_template!(rx_callback, _stack, {
+fn rx_callback(_: &mut InterruptStack) {
     RX_FLAG.store(true, Ordering::SeqCst);
-});
+}
 
-interrupt_handler_template!(tx_callback, _stack, {
+fn tx_callback(_: &mut InterruptStack) {
     TX_FLAG.store(true, Ordering::SeqCst);
-});
+}
 
-interrupt_handler_template!(config_callback, _stack, {
+fn config_callback(_: &mut InterruptStack) {
     CONFIG_FLAG.store(true, Ordering::SeqCst);
-});
+}

src/devices/vsock/src/transport/vmcall.rs

@@ -16,10 +16,9 @@ use super::Result;
 use alloc::boxed::Box;
 use alloc::collections::BTreeMap;
 use alloc::vec::Vec;
-use core::arch::asm;
 use core::convert::TryInto;
 use core::sync::atomic::{AtomicBool, Ordering};
-use td_payload::{eoi, interrupt_handler_template};
+use td_payload::arch::idt::InterruptStack;
 use td_shim_interface::td_uefi_pi::pi::guid;
 use tdx_tdcall::tdx;
 
@@ -54,7 +53,7 @@ impl VmcallVsock {
         dma_allocator: Box<dyn VsockDmaPageAllocator>,
         timer: Box<dyn VsockTimeout>,
     ) -> Result<Self> {
-        register_callback(VMCALL_VECTOR, vmcall_notification);
+        register_callback(VMCALL_VECTOR, vmcall_notification)?;
 
         Ok(Self {
             mid,
@@ -406,9 +405,9 @@ impl<'a> Response<'a> {
     }
 }
 
-interrupt_handler_template!(vmcall_notification, _stack, {
+fn vmcall_notification(_: &mut InterruptStack) {
     VMCALL_FLAG.store(true, Ordering::SeqCst);
-});
+}
 
 fn speculation_barrier() {
     unsafe { core::arch::x86_64::_mm_lfence() };

src/migtd/src/driver/timer.rs

@@ -4,8 +4,7 @@
 
 use core::sync::atomic::{AtomicBool, Ordering};
 use td_payload::arch::apic::*;
-use td_payload::arch::idt::register;
-use td_payload::interrupt_handler_template;
+use td_payload::arch::idt::{register_interrupt_callback, InterruptCallback, InterruptStack};
 
 /// A simple apic timer notification handler used to handle the
 /// time out events
@@ -15,9 +14,9 @@ static TIMEOUT_FLAG: AtomicBool = AtomicBool::new(false);
 const TIMEOUT_VECTOR: u8 = 33;
 const CPUID_TSC_DEADLINE_BIT: u32 = 1 << 24;
 
-interrupt_handler_template!(timer, _stack, {
+fn timer_handler(_stack: &mut InterruptStack) {
     TIMEOUT_FLAG.store(true, Ordering::SeqCst);
-});
+}
 
 pub fn init_timer() {
     let cpuid = unsafe { core::arch::x86_64::__cpuid_count(0x1, 0) };
@@ -58,5 +57,8 @@ fn apic_timer_lvtt_setup(vector: u8) {
 
 fn set_timer_notification(vector: u8) {
     // Setup interrupt handler
-    register(vector, timer);
+    if register_interrupt_callback(vector as usize, InterruptCallback::new(timer_handler)).is_err()
+    {
+        panic!("Failed to set interrupt callback for timer");
+    }
 }

src/migtd/src/lib.rs

@@ -46,7 +46,7 @@ pub extern "C" fn _start(hob: u64, payload: u64) -> ! {
         shadow_stack_size: DEFAULT_SHADOW_STACK_SIZE,
     };
 
-    arch::init::pre_init(hob as u64, &layout);
+    arch::init::pre_init(hob as u64, &layout, true);
 
     // Init internal heap
     #[cfg(not(feature = "test_disable_ra_and_accept_all"))]

src/migtd/src/migration/event.rs

@@ -4,18 +4,24 @@
 
 use core::sync::atomic::{AtomicBool, Ordering};
 use td_payload::arch::apic::*;
-use td_payload::arch::idt::register;
-use td_payload::interrupt_handler_template;
+use td_payload::arch::idt::{register_interrupt_callback, InterruptCallback, InterruptStack};
 
 pub const VMCALL_SERVICE_VECTOR: u8 = 0x50;
 pub static VMCALL_SERVICE_FLAG: AtomicBool = AtomicBool::new(false);
 
-interrupt_handler_template!(vmcall_service_callback, _stack, {
+fn vmcall_service_callback(_stack: &mut InterruptStack) {
     VMCALL_SERVICE_FLAG.store(true, Ordering::SeqCst);
-});
+}
 
 pub fn register_callback() {
-    register(VMCALL_SERVICE_VECTOR, vmcall_service_callback);
+    if register_interrupt_callback(
+        VMCALL_SERVICE_VECTOR as usize,
+        InterruptCallback::new(vmcall_service_callback),
+    )
+    .is_err()
+    {
+        panic!("Failed to set interrupt callback for VMCALL_SERVICE");
+    }
 }
 
 pub fn wait_for_event(event_flag: &AtomicBool) {