vcpu: add support for kicking out of KVM_RUN

acatangiu · alxiord · commit b359722c23ec · 2019-12-04T14:25:30.000-08:00
This is done by making use of the kvm `immediate_exit` flag.

The vcpu is signaled and the `immediate_exit` flag is set in the
signal handler. This leads to race-free interruption of KVM_RUN.

In order to access the `immediate_exit` flag for a particular vcpu
from a global context such as the signal handler, we use the TLS
of that vcpu thread to gain access to the vcpu's data.

This commit tries to keep the flow as safe as possible, and
in that spirit, the following vcpu methods are defined:
- `init_thread_local_data(&amp;mut self)` - associates the `self` vcpu
   to the current thread.
- `reset_thread_local_data(&amp;mut self)` - deassociates the `self`
   vcpu from the current thread.
- `Vcpu::run_on_thread_local&lt;F&gt;(fn: F)` - runs `fn` passing it
   a reference to the vcpu in TLS of local thread. This model
   enforces a limited lifetime (lifetime of `fn`) for the unsafe
   vcpu reference.
- `Vcpu::register_kick_signal_handler()` - Registers a signal
   handler that sets kvm `immediate_exit` using
  `run_on_thread_local()` in order to kick the vcpu running on
   the current thread.

Signed-off-by: Adrian Catangiu &lt;acatan@amazon.com&gt;
diff --git a/src/vmm/src/lib.rs b/src/vmm/src/lib.rs
@@ -914,6 +914,8 @@ impl Vmm {
             "The number of vCPU fds is corrupted!"
         );
 
+        Vcpu::register_kick_signal_handler();
+
         self.vcpus_handles.reserve(vcpu_count as usize);
 
         let vcpus_thread_barrier = Arc::new(Barrier::new((vcpu_count + 1) as usize));
diff --git a/src/vmm/src/vstate.rs b/src/vmm/src/vstate.rs
@@ -5,8 +5,11 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the THIRD-PARTY file.
 
+use libc::{c_int, c_void, siginfo_t};
+use std::cell::Cell;
 use std::io;
 use std::result;
+use std::sync::atomic::{fence, Ordering};
 use std::sync::{Arc, Barrier};
 
 use super::TimestampUs;
@@ -23,6 +26,7 @@ use kvm_ioctls::*;
 use logger::{LogOption, Metric, LOGGER, METRICS};
 use memory_model::{Address, GuestAddress, GuestMemory, GuestMemoryError};
 use utils::eventfd::EventFd;
+use utils::signal::{register_signal_handler, SignalHandler};
 #[cfg(target_arch = "x86_64")]
 use vmm_config::machine_config::{CpuFeaturesTemplate, VmConfig};
 
@@ -34,6 +38,8 @@ const MAGIC_IOPORT_SIGNAL_GUEST_BOOT_COMPLETE: u64 = 0x03f0;
 const MAGIC_IOPORT_SIGNAL_GUEST_BOOT_COMPLETE: u64 = 0x40000000;
 const MAGIC_VALUE_SIGNAL_GUEST_BOOT_COMPLETE: u8 = 123;
 
+pub(crate) const VCPU_RTSIG_OFFSET: i32 = 0;
+
 /// Errors associated with the wrappers over KVM ioctls.
 #[derive(Debug)]
 pub enum Error {
@@ -86,6 +92,10 @@ pub enum Error {
     Irq(kvm_ioctls::Error),
     /// Cannot spawn a new vCPU thread.
     VcpuSpawn(io::Error),
+    /// Cannot clean init vcpu TLS.
+    VcpuTlsInit,
+    /// Vcpu not present in TLS.
+    VcpuTlsNotPresent,
     /// Unexpected KVM_RUN exit reason
     VcpuUnhandledKvmExit,
     #[cfg(target_arch = "aarch64")]
@@ -267,6 +277,9 @@ impl Vm {
     }
 }
 
+// Using this for easier explicit type-casting to help IDEs interpret the code.
+type VcpuCell = Cell<Option<*const Vcpu>>;
+
 /// A wrapper around creating and using a kvm-based VCPU.
 pub struct Vcpu {
     #[cfg(target_arch = "x86_64")]
@@ -282,6 +295,95 @@ pub struct Vcpu {
 }
 
 impl Vcpu {
+    thread_local!(static TLS_VCPU_PTR: VcpuCell = Cell::new(None));
+
+    /// Associates `self` with the current thread.
+    ///
+    /// It is a prerequisite to successfully run `init_thread_local_data()` before using
+    /// `run_on_thread_local()` on the current thread.
+    /// This function will return an error if there already is a `Vcpu` present in the TLS.
+    fn init_thread_local_data(&mut self) -> Result<()> {
+        Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| {
+            if cell.get().is_some() {
+                return Err(Error::VcpuTlsInit);
+            }
+            cell.set(Some(self as *const Vcpu));
+            Ok(())
+        })
+    }
+
+    /// Deassociates `self` from the current thread.
+    ///
+    /// Should be called if the current `self` had called `init_thread_local_data()` and
+    /// now needs to move to a different thread.
+    ///
+    /// Fails if `self` was not previously associated with the current thread.
+    fn reset_thread_local_data(&mut self) -> Result<()> {
+        // Best-effort to clean up TLS. If the `Vcpu` was moved to another thread
+        // _before_ running this, then there is nothing we can do.
+        Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| {
+            if let Some(vcpu_ptr) = cell.get() {
+                if vcpu_ptr == self as *const Vcpu {
+                    Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| cell.take());
+                    return Ok(());
+                }
+            }
+            Err(Error::VcpuTlsNotPresent)
+        })
+    }
+
+    /// Runs `func` for the `Vcpu` associated with the current thread.
+    ///
+    /// It requires that `init_thread_local_data()` was run on this thread.
+    ///
+    /// Fails if there is no `Vcpu` associated with the current thread.
+    ///
+    /// # Safety
+    ///
+    /// This is marked unsafe as it allows temporary aliasing through
+    /// dereferencing from pointer an already borrowed `Vcpu`.
+    unsafe fn run_on_thread_local<F>(func: F) -> Result<()>
+    where
+        F: FnOnce(&Vcpu),
+    {
+        Self::TLS_VCPU_PTR.with(|cell: &VcpuCell| {
+            if let Some(vcpu_ptr) = cell.get() {
+                // Dereferencing here is safe since `TLS_VCPU_PTR` is populated/non-empty,
+                // and it is being cleared on `Vcpu::drop` so there is no dangling pointer.
+                let vcpu_ref: &Vcpu = &*vcpu_ptr;
+                func(vcpu_ref);
+                Ok(())
+            } else {
+                Err(Error::VcpuTlsNotPresent)
+            }
+        })
+    }
+
+    /// Registers a signal handler which makes use of TLS and kvm immediate exit to
+    /// kick the vcpu running on the current thread, if there is one.
+    pub fn register_kick_signal_handler() {
+        extern "C" fn handle_signal(_: c_int, _: *mut siginfo_t, _: *mut c_void) {
+            // This is safe because it's temporarily aliasing the `Vcpu` object, but we are
+            // only reading `vcpu.fd` which does not change for the lifetime of the `Vcpu`.
+            unsafe {
+                let _ = Vcpu::run_on_thread_local(|vcpu| {
+                    vcpu.fd.set_kvm_immediate_exit(1);
+                    fence(Ordering::Release);
+                });
+            }
+        }
+        unsafe {
+            // This uses an async signal safe handler to kill the vcpu handles.
+            register_signal_handler(
+                VCPU_RTSIG_OFFSET,
+                SignalHandler::Siginfo(handle_signal),
+                true,
+                libc::SA_SIGINFO,
+            )
+            .expect("Failed to register vcpu signal handler");
+        }
+    }
+
     /// Constructs a new VCPU for `vm`.
     ///
     /// # Arguments
@@ -439,29 +541,32 @@ impl Vcpu {
         }
     }
 
-    fn run_emulation(&mut self) -> Result<()> {
+    /// Runs the vCPU in KVM context and handles the kvm exit reason.
+    ///
+    /// Returns error or enum specifying whether emulation was handled or interrupted.
+    fn run_emulation(&mut self) -> Result<VcpuEmulation> {
         match self.fd.run() {
             Ok(run) => match run {
                 #[cfg(target_arch = "x86_64")]
                 VcpuExit::IoIn(addr, data) => {
                     self.io_bus.read(u64::from(addr), data);
                     METRICS.vcpu.exit_io_in.inc();
-                    Ok(())
+                    Ok(VcpuEmulation::Handled)
                 }
                 #[cfg(target_arch = "x86_64")]
                 VcpuExit::IoOut(addr, data) => {
                     self.check_boot_complete_signal(u64::from(addr), data);
 
                     self.io_bus.write(u64::from(addr), data);
                     METRICS.vcpu.exit_io_out.inc();
-                    Ok(())
+                    Ok(VcpuEmulation::Handled)
                 }
                 VcpuExit::MmioRead(addr, data) => {
                     if let Some(ref mmio_bus) = self.mmio_bus {
                         mmio_bus.read(addr, data);
                         METRICS.vcpu.exit_mmio_read.inc();
                     }
-                    Ok(())
+                    Ok(VcpuEmulation::Handled)
                 }
                 VcpuExit::MmioWrite(addr, data) => {
                     if let Some(ref mmio_bus) = self.mmio_bus {
@@ -471,7 +576,7 @@ impl Vcpu {
                         mmio_bus.write(addr, data);
                         METRICS.vcpu.exit_mmio_write.inc();
                     }
-                    Ok(())
+                    Ok(VcpuEmulation::Handled)
                 }
                 VcpuExit::Hlt => {
                     info!("Received KVM_EXIT_HLT signal");
@@ -505,8 +610,12 @@ impl Vcpu {
             // error in our code in which case it is better to panic.
             Err(ref e) => {
                 match e.errno() {
-                    // Why do we check for these if we only return EINVAL?
-                    libc::EAGAIN | libc::EINTR => Ok(()),
+                    libc::EAGAIN => Ok(VcpuEmulation::Handled),
+                    libc::EINTR => {
+                        self.fd.set_kvm_immediate_exit(0);
+                        // Notify that this KVM_RUN was interrupted.
+                        Ok(VcpuEmulation::Interrupted)
+                    }
                     _ => {
                         METRICS.vcpu.failures.inc();
                         error!("Failure during vcpu run: {}", e);
@@ -528,6 +637,9 @@ impl Vcpu {
         seccomp_level: u32,
         vcpu_exit_evt: EventFd,
     ) {
+        self.init_thread_local_data()
+            .expect("Cannot cleanly initialize vcpu TLS.");
+
         // Load seccomp filters for this vCPU thread.
         // Execution panics if filters cannot be loaded, use --seccomp-level=0 if skipping filters
         // altogether is the desired behaviour.
@@ -550,12 +662,24 @@ impl Vcpu {
     }
 }
 
+impl Drop for Vcpu {
+    fn drop(&mut self) {
+        let _ = self.reset_thread_local_data();
+    }
+}
+
+enum VcpuEmulation {
+    Handled,
+    Interrupted,
+}
+
 #[cfg(test)]
 mod tests {
     use std::fs::File;
 
     use super::super::devices;
     use super::*;
+    use utils::signal::Killable;
 
     // Auxiliary function being used throughout the tests.
     fn setup_vcpu() -> (Vm, Vcpu) {
@@ -740,4 +864,85 @@ mod tests {
         assert_eq!(m1.dev(), m2.dev());
         assert_eq!(m1.ino(), m2.ino());
     }
+
+    #[test]
+    fn test_vcpu_tls() {
+        let (_, mut vcpu) = setup_vcpu();
+
+        // Running on the TLS vcpu should fail before we actually initialize it.
+        unsafe {
+            assert!(Vcpu::run_on_thread_local(|_| ()).is_err());
+        }
+
+        // Initialize vcpu TLS.
+        vcpu.init_thread_local_data().unwrap();
+
+        // Validate TLS vcpu is the local vcpu by changing the `id` then validating against
+        // the one in TLS.
+        vcpu.id = 12;
+        unsafe {
+            assert!(Vcpu::run_on_thread_local(|v| assert_eq!(v.id, 12)).is_ok());
+        }
+
+        // Reset vcpu TLS.
+        assert!(vcpu.reset_thread_local_data().is_ok());
+
+        // Running on the TLS vcpu after TLS reset should fail.
+        unsafe {
+            assert!(Vcpu::run_on_thread_local(|_| ()).is_err());
+        }
+
+        // Second reset should return error.
+        assert!(vcpu.reset_thread_local_data().is_err());
+    }
+
+    #[test]
+    fn test_invalid_tls() {
+        let (_, mut vcpu) = setup_vcpu();
+        // Initialize vcpu TLS.
+        vcpu.init_thread_local_data().unwrap();
+        // Trying to initialize non-empty TLS should error.
+        vcpu.init_thread_local_data().unwrap_err();
+    }
+
+    #[test]
+    fn test_vcpu_kick() {
+        Vcpu::register_kick_signal_handler();
+        let (vm, mut vcpu) = setup_vcpu();
+
+        let kvm_run =
+            KvmRunWrapper::mmap_from_fd(&vcpu.fd, vm.fd.run_size()).expect("cannot mmap kvm-run");
+        let success = Arc::new(std::sync::atomic::AtomicBool::new(false));
+        let vcpu_success = success.clone();
+        let barrier = Arc::new(Barrier::new(2));
+        let vcpu_barrier = barrier.clone();
+        // Start Vcpu thread which will be kicked with a signal.
+        let handle = std::thread::Builder::new()
+            .name("test_vcpu_kick".to_string())
+            .spawn(move || {
+                vcpu.init_thread_local_data().unwrap();
+                // Notify TLS was populated.
+                vcpu_barrier.wait();
+                // Loop for max 1 second to check if the signal handler has run.
+                for _ in 0..10 {
+                    if kvm_run.as_mut_ref().immediate_exit == 1 {
+                        // Signal handler has run and set immediate_exit to 1.
+                        vcpu_success.store(true, Ordering::Release);
+                        break;
+                    }
+                    std::thread::sleep(std::time::Duration::from_millis(100));
+                }
+            })
+            .expect("cannot start thread");
+
+        // Wait for the vcpu to initialize its TLS.
+        barrier.wait();
+        // Kick the Vcpu using the custom signal.
+        handle
+            .kill(VCPU_RTSIG_OFFSET)
+            .expect("failed to signal thread");
+        handle.join().expect("failed to join thread");
+        // Verify that the Vcpu saw its kvm immediate-exit as set.
+        assert!(success.load(Ordering::Acquire));
+    }
 }
