refactor(its): cleanup ITS save/restore logic

bchalios · bchalios · commit e5323863f10b · 2025-06-25T12:24:06.000+02:00
We should squash that to previous commit

Signed-off-by: Babis Chalios &lt;bchalios@amazon.es&gt;
diff --git a/src/vmm/src/arch/aarch64/gic/gicv3/mod.rs b/src/vmm/src/arch/aarch64/gic/gicv3/mod.rs
@@ -1,7 +1,7 @@
 // Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-mod regs;
+pub mod regs;
 
 use kvm_ioctls::{DeviceFd, VmFd};
 
@@ -234,7 +234,7 @@ impl GICv3 {
 /// RDIST pending tables into guest RAM.
 ///
 /// The tables get flushed to guest RAM whenever the VM gets stopped.
-fn save_pending_tables(gic_device: &DeviceFd, its_device: &DeviceFd) -> Result<(), GicError> {
+fn save_pending_tables(gic_device: &DeviceFd) -> Result<(), GicError> {
     let init_gic_attr = kvm_bindings::kvm_device_attr {
         group: kvm_bindings::KVM_DEV_ARM_VGIC_GRP_CTRL,
         attr: u64::from(kvm_bindings::KVM_DEV_ARM_VGIC_SAVE_PENDING_TABLES),
@@ -243,18 +243,7 @@ fn save_pending_tables(gic_device: &DeviceFd, its_device: &DeviceFd) -> Result<(
     };
     gic_device.set_device_attr(&init_gic_attr).map_err(|err| {
         GicError::DeviceAttribute(err, true, kvm_bindings::KVM_DEV_ARM_VGIC_GRP_CTRL)
-    })?;
-    let save_its_tables_attr = kvm_bindings::kvm_device_attr {
-        group: kvm_bindings::KVM_DEV_ARM_VGIC_GRP_CTRL,
-        attr: u64::from(kvm_bindings::KVM_DEV_ARM_ITS_SAVE_TABLES),
-        addr: 0,
-        flags: 0,
-    };
-    its_device
-        .set_device_attr(&save_its_tables_attr)
-        .map_err(|err| {
-            GicError::DeviceAttribute(err, true, kvm_bindings::KVM_DEV_ARM_VGIC_GRP_CTRL)
-        })
+    })
 }
 
 #[cfg(test)]
@@ -272,11 +261,11 @@ mod tests {
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let gic = create_gic(&vm, 1, Some(GICVersion::GICV3)).expect("Cannot create gic");
-        save_pending_tables(gic.device_fd(), gic.its_fd().unwrap()).unwrap();
+        save_pending_tables(gic.device_fd()).unwrap();
 
         unsafe { libc::close(gic.device_fd().as_raw_fd()) };
 
-        let res = save_pending_tables(gic.device_fd(), gic.its_fd().unwrap());
+        let res = save_pending_tables(gic.device_fd());
         assert_eq!(
             format!("{:?}", res.unwrap_err()),
             "DeviceAttribute(Error(9), true, 4)"
diff --git a/src/vmm/src/arch/aarch64/gic/gicv3/regs/its_regs.rs b/src/vmm/src/arch/aarch64/gic/gicv3/regs/its_regs.rs
@@ -0,0 +1,124 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use kvm_bindings::{
+    KVM_DEV_ARM_ITS_RESTORE_TABLES, KVM_DEV_ARM_ITS_SAVE_TABLES, KVM_DEV_ARM_VGIC_GRP_CTRL,
+    KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
+};
+use kvm_ioctls::DeviceFd;
+use serde::{Deserialize, Serialize};
+
+use crate::arch::aarch64::gic::GicError;
+
+// ITS registers that we want to preserve across snapshots
+const GITS_CTLR: u32 = 0x0000;
+const GITS_IIDR: u32 = 0x0004;
+const GITS_CBASER: u32 = 0x0080;
+const GITS_CWRITER: u32 = 0x0088;
+const GITS_CREADR: u32 = 0x0090;
+const GITS_BASER: u32 = 0x0100;
+
+fn set_device_attribute(
+    its_device: &DeviceFd,
+    group: u32,
+    attr: u32,
+    val: u64,
+) -> Result<(), GicError> {
+    let gicv3_its_attr = kvm_bindings::kvm_device_attr {
+        group,
+        attr: attr as u64,
+        addr: &val as *const u64 as u64,
+        flags: 0,
+    };
+
+    its_device
+        .set_device_attr(&gicv3_its_attr)
+        .map_err(|err| GicError::DeviceAttribute(err, true, group))
+}
+
+fn get_device_attribute(its_device: &DeviceFd, group: u32, attr: u32) -> Result<u64, GicError> {
+    let mut val = 0;
+
+    let mut gicv3_its_attr = kvm_bindings::kvm_device_attr {
+        group,
+        attr: attr as u64,
+        addr: &mut val as *mut u64 as u64,
+        flags: 0,
+    };
+
+    // SAFETY: gicv3_its_attr.addr is safe to write to.
+    unsafe { its_device.get_device_attr(&mut gicv3_its_attr) }
+        .map_err(|err| GicError::DeviceAttribute(err, false, group))?;
+
+    Ok(val)
+}
+
+fn its_read_register(its_fd: &DeviceFd, attr: u32) -> Result<u64, GicError> {
+    get_device_attribute(its_fd, KVM_DEV_ARM_VGIC_GRP_ITS_REGS, attr)
+}
+
+fn its_set_register(its_fd: &DeviceFd, attr: u32, val: u64) -> Result<(), GicError> {
+    set_device_attribute(its_fd, KVM_DEV_ARM_VGIC_GRP_ITS_REGS, attr, val)
+}
+
+pub fn its_save_tables(its_fd: &DeviceFd) -> Result<(), GicError> {
+    set_device_attribute(
+        its_fd,
+        KVM_DEV_ARM_VGIC_GRP_CTRL,
+        KVM_DEV_ARM_ITS_SAVE_TABLES,
+        0,
+    )
+}
+
+pub fn its_restore_tables(its_fd: &DeviceFd) -> Result<(), GicError> {
+    set_device_attribute(
+        its_fd,
+        KVM_DEV_ARM_VGIC_GRP_CTRL,
+        KVM_DEV_ARM_ITS_RESTORE_TABLES,
+        0,
+    )
+}
+
+/// ITS registers that we save/restore during snapshot
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub struct ItsRegisterState {
+    iidr: u64,
+    cbaser: u64,
+    creadr: u64,
+    cwriter: u64,
+    baser: [u64; 8],
+    ctlr: u64,
+}
+
+impl ItsRegisterState {
+    /// Save ITS state
+    pub fn save(its_fd: &DeviceFd) -> Result<Self, GicError> {
+        let mut state = ItsRegisterState::default();
+
+        for i in 0..8 {
+            state.baser[i as usize] = its_read_register(its_fd, GITS_BASER + i * 8)?;
+        }
+        state.ctlr = its_read_register(its_fd, GITS_CTLR)?;
+        state.cbaser = its_read_register(its_fd, GITS_CBASER)?;
+        state.creadr = its_read_register(its_fd, GITS_CREADR)?;
+        state.cwriter = its_read_register(its_fd, GITS_CWRITER)?;
+        state.iidr = its_read_register(its_fd, GITS_IIDR)?;
+
+        Ok(state)
+    }
+
+    /// Restore ITS state
+    pub fn restore(&self, its_fd: &DeviceFd) -> Result<(), GicError> {
+        its_set_register(its_fd, GITS_IIDR, self.iidr)?;
+        its_set_register(its_fd, GITS_CBASER, self.cbaser)?;
+        its_set_register(its_fd, GITS_CREADR, self.creadr)?;
+        its_set_register(its_fd, GITS_CWRITER, self.cwriter)?;
+        for i in 0..8 {
+            its_set_register(its_fd, GITS_BASER + i * 8, self.baser[i as usize])?;
+        }
+
+        // We need to restore saved ITS tables before restoring GITS_CTLR
+        its_restore_tables(its_fd)?;
+        its_set_register(its_fd, GITS_CTLR, self.ctlr)
+    }
+}
diff --git a/src/vmm/src/arch/aarch64/gic/gicv3/regs/mod.rs b/src/vmm/src/arch/aarch64/gic/gicv3/regs/mod.rs
@@ -3,54 +3,14 @@
 
 mod dist_regs;
 mod icc_regs;
+pub mod its_regs;
 mod redist_regs;
 
 use kvm_ioctls::DeviceFd;
 
 use crate::arch::aarch64::gic::GicError;
 use crate::arch::aarch64::gic::regs::{GicState, GicVcpuState};
-
-const GITS_CTLR: u32 = 0x0000;
-const GITS_IIDR: u32 = 0x0004;
-const GITS_CBASER: u32 = 0x0080;
-const GITS_CWRITER: u32 = 0x0088;
-const GITS_CREADR: u32 = 0x0090;
-const GITS_BASER: u32 = 0x0100;
-
-pub fn gicv3_its_attr_set(
-    its_device: &DeviceFd,
-    group: u32,
-    attr: u32,
-    val: u64,
-) -> Result<(), GicError> {
-    let gicv3_its_attr = kvm_bindings::kvm_device_attr {
-        group,
-        attr: attr as u64,
-        addr: &val as *const u64 as u64,
-        flags: 0,
-    };
-
-    its_device
-        .set_device_attr(&gicv3_its_attr)
-        .map_err(|err| GicError::DeviceAttribute(err, true, group))
-}
-
-pub fn gicv3_its_attr_get(its_device: &DeviceFd, group: u32, attr: u32) -> Result<u64, GicError> {
-    let mut val = 0;
-
-    let mut gicv3_its_attr = kvm_bindings::kvm_device_attr {
-        group,
-        attr: attr as u64,
-        addr: &mut val as *mut u64 as u64,
-        flags: 0,
-    };
-
-    // SAFETY: gicv3_its_attr.addr is safe to write to.
-    unsafe { its_device.get_device_attr(&mut gicv3_its_attr) }
-        .map_err(|err| GicError::DeviceAttribute(err, false, group))?;
-
-    Ok(val)
-}
+use its_regs::{ItsRegisterState, its_save_tables};
 
 /// Save the state of the GIC device.
 pub fn save_state(
@@ -59,7 +19,9 @@ pub fn save_state(
     mpidrs: &[u64],
 ) -> Result<GicState, GicError> {
     // Flush redistributors pending tables to guest RAM.
-    super::save_pending_tables(gic_device, its_device)?;
+    super::save_pending_tables(gic_device)?;
+    // Flush ITS tables into guest memory.
+    its_save_tables(its_device)?;
 
     let mut vcpu_states = Vec::with_capacity(mpidrs.len());
     for mpidr in mpidrs {
@@ -69,54 +31,12 @@ pub fn save_state(
         })
     }
 
-    let mut its_baser: [u64; 8] = [0; 8];
-    for i in 0..8 {
-        its_baser[i as usize] = gicv3_its_attr_get(
-            its_device,
-            kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-            GITS_BASER + i * 8,
-        )?;
-    }
-
-    let its_ctlr = gicv3_its_attr_get(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CTLR,
-    )?;
-
-    let its_cbaser = gicv3_its_attr_get(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CBASER,
-    )?;
-
-    let its_creadr = gicv3_its_attr_get(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CREADR,
-    )?;
-
-    let its_cwriter = gicv3_its_attr_get(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CWRITER,
-    )?;
-
-    let its_iidr = gicv3_its_attr_get(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_IIDR,
-    )?;
+    let its_state = ItsRegisterState::save(its_device)?;
 
     Ok(GicState {
         dist: dist_regs::get_dist_regs(gic_device)?,
         gic_vcpu_states: vcpu_states,
-        its_ctlr,
-        its_cbaser,
-        its_creadr,
-        its_cwriter,
-        its_iidr,
-        its_baser,
+        its_state: Some(its_state),
     })
 }
 
@@ -137,59 +57,9 @@ pub fn restore_state(
         icc_regs::set_icc_regs(gic_device, *mpidr, &vcpu_state.icc)?;
     }
 
-    // Restore ITS registers
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_IIDR,
-        state.its_iidr,
-    )?;
-
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CBASER,
-        state.its_cbaser,
-    )?;
-
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CREADR,
-        state.its_creadr,
-    )?;
-
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CWRITER,
-        state.its_cwriter,
-    )?;
-
-    for i in 0..8 {
-        gicv3_its_attr_set(
-            its_device,
-            kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-            GITS_BASER + i * 8,
-            state.its_baser[i as usize],
-        )?;
-    }
-
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_CTRL,
-        kvm_bindings::KVM_DEV_ARM_ITS_RESTORE_TABLES,
-        0,
-    )?;
-
-    gicv3_its_attr_set(
-        its_device,
-        kvm_bindings::KVM_DEV_ARM_VGIC_GRP_ITS_REGS,
-        GITS_CTLR,
-        state.its_ctlr,
-    )?;
-
-    Ok(())
+    // Safe to unwrap here, as we know we support an ITS device, so `its_state.is_some()` is always
+    // `true`.
+    state.its_state.as_ref().unwrap().restore(its_device)
 }
 
 #[cfg(test)]
diff --git a/src/vmm/src/arch/aarch64/gic/regs.rs b/src/vmm/src/arch/aarch64/gic/regs.rs
@@ -10,6 +10,7 @@ use kvm_ioctls::DeviceFd;
 use serde::{Deserialize, Serialize};
 
 use crate::arch::aarch64::gic::GicError;
+use crate::arch::aarch64::gic::gicv3::regs::its_regs::ItsRegisterState;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct GicRegState<T> {
@@ -30,18 +31,8 @@ pub struct GicState {
     pub dist: Vec<GicRegState<u32>>,
     /// The state of the vcpu interfaces.
     pub gic_vcpu_states: Vec<GicVcpuState>,
-    /// ITS control register
-    pub its_ctlr: u64,
-    /// ITS IID register
-    pub its_iidr: u64,
-    /// ITS CBASE register
-    pub its_cbaser: u64,
-    /// ITS CWRITE register
-    pub its_cwriter: u64,
-    /// ITS CREAD register
-    pub its_creadr: u64,
-    /// ITS BASE registers
-    pub its_baser: [u64; 8],
+    /// The state of the ITS device. Only present with GICv3.
+    pub its_state: Option<ItsRegisterState>,
 }
 
 /// Structure used for serializing the state of the GIC registers for a specific vCPU.
