vmm: fix immediate_exit handling for PIO/MMIO code paths

To finish the emulation of PIO or MMIO code paths [0], we need to get
back into KVM_RUN. In order to not lose any signal that we received,
we must preserve the original `immediate_exit` flag in case it was
already set by the signal handler.

[0] https://elixir.bootlin.com/linux/v6.12/source/arch/x86/kvm/x86.c#L11555

Author: phip1611

vmm/src/cpu.rs

@@ -1121,16 +1121,9 @@ impl CpuManager {
                         }
 
                         let kvm_run = KVM_RUN.get();
-                        assert!(!kvm_run.is_null(), "kvm_run should have been mapped as part of vCPU setup");
-                        // SAFETY: We mapped the whole structure.
-                        let ptr_kvm_run_immediate_exit = unsafe { kvm_run.cast::<u8>().add(1) };
-                        // SAFETY: We know the mapping is valid.
-                        unsafe {
-                            core::ptr::write_volatile(ptr_kvm_run_immediate_exit, 1);
-                        }
-
-                        // TODO ::Release?
-                        std::sync::atomic::fence(Ordering::SeqCst);
+                        let kvm_run = unsafe {
+                            kvm_run.as_mut().expect("kvm_run should have been mapped as part of vCPU setup") };
+                        kvm_run.immediate_exit = 1;
                     }
                     // This uses an async signal safe handler to kill the vcpu handles.
                     register_signal_handler(SIGRTMIN(), handle_signal)
@@ -1170,12 +1163,17 @@ impl CpuManager {
 
                                 #[cfg(feature = "kvm")]
                                 if matches!(hypervisor_type, HypervisorType::Kvm) {
+                                    let old_immediate_exit_flag = {
+                                        let kvm_run: *mut kvm_run = KVM_RUN.get();
+                                        let kvm_run = unsafe { kvm_run.as_mut().unwrap() };
+                                        kvm_run.immediate_exit != 0
+                                    };
                                     vcpu.lock().as_ref().unwrap().vcpu.set_immediate_exit(true);
                                     if !matches!(vcpu.lock().unwrap().run(), Ok(VmExit::Ignore)) {
                                         error!("Unexpected VM exit on \"immediate_exit\" run");
                                         break;
                                     }
-                                    vcpu.lock().as_ref().unwrap().vcpu.set_immediate_exit(false);
+                                    vcpu.lock().as_ref().unwrap().vcpu.set_immediate_exit(old_immediate_exit_flag);
                                 }
 
                                 vcpu_run_interrupted.store(true, Ordering::SeqCst);