Snapshot editor

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+- [#4045](https://github.com/firecracker-microvm/firecracker/pull/4045)
+  Added `snapshot-editor` tool for modifications of snapshot files.
 - [#3967](https://github.com/firecracker-microvm/firecracker/pull/3967/):
   Added new fields to the custom CPU templates. (aarch64 only) `vcpu_features`
   field allows modifications of vCPU features enabled during vCPU

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["src/cpu-template-helper", "src/firecracker", "src/jailer", "src/rebase-snap", "src/seccompiler"]
+members = ["src/cpu-template-helper", "src/firecracker", "src/jailer", "src/rebase-snap", "src/seccompiler", "src/snapshot-editor"]
 default-members = ["src/firecracker"]
 resolver = "2"
 

docs/snapshotting/snapshot-editor.md

@@ -0,0 +1,91 @@
+# Snapshot editor
+
+The `snapshot-editor` is a program for modification of Firecracker snapshots.
+
+## Prior knowledge
+
+Firecracker snapshot consists of 2 files:
+
+- `vmstate` file: file with Firecracker internal data such as vcpu states,
+  devices states etc.
+- `memory` file: file with guest memory.
+
+`snapshot-editor` (currently) allows to modify only `vmstate` files only
+on aarch64 platform.
+
+## Usage
+
+### `edit-vmstate` command
+
+#### `remove-regs` subcommand (aarch64 only)
+
+This command is used to remove specified registers from vcpu states inside
+vmstate snapshot file.
+
+Arguments:
+
+- `VMSTATE_PATH` - path to the `vmstate` file
+- `OUTPUT_PATH` - path to the file where the output will be placed
+- `[REGS]` - set of u32 values representing registers ids as they are defined
+  in KVM. Can be both in decimal and in hex formats.
+
+Usage:
+
+```bash
+snapshot-editor edit-vmstate remove-regs \
+    --vmstate-path <VMSTATE_PATH> \
+    --output-path <OUTPUT_PATH> \
+    [REGS]...
+```
+
+Example:
+
+```bash
+./snapshot-editor edit-vmstate remove-regs \
+    --vmstate-path ./vmstate_file \
+    --output-path ./new_vmstate_file \
+    0x1 0x2
+```
+
+### `info-vmstate` command
+
+#### `version` subcommand
+
+This command is used to print version of the provided
+vmstate file.
+
+Arguments:
+
+- `VMSTATE_PATH` - path to the `vmstate` file
+
+Usage:
+
+```bash
+snapshot-editor info-vmstate version --vmstate-path <VMSTATE_PATH>
+```
+
+Example:
+
+```bash
+./snapshot-editor info-vmstate version --vmstate-path ./vmstate_file
+```
+
+#### `vcpu-states` subcommand (aarch64 only)
+
+This command is used to print the vCPU states inside vmstate snapshot file.
+
+Arguments:
+
+- `VMSTATE_PATH` - path to the `vmstate` file
+
+Usage:
+
+```bash
+snapshot-editor info-vmstate vcpu-states --vmstate-path <VMSTATE_PATH>
+```
+
+Example:
+
+```bash
+./snapshot-editor info-vmstate vcpu-states --vmstate-path ./vmstate_file
+```

src/snapshot-editor/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "snapshot-editor"
+version = "1.5.0"
+authors = ["Amazon Firecracker team <[email protected]>"]
+edition = "2021"
+build = "../../build.rs"
+license = "Apache-2.0"
+
+[dependencies]
+clap = { version = "4.3.19", features = ["derive", "string"] }
+clap-num = "1.0.2"
+libc = "0.2.147"
+snapshot = { path = "../snapshot" }
+thiserror = "1.0.44"
+vmm = { path = "../vmm" }

src/snapshot-editor/src/edit_vmstate.rs

@@ -0,0 +1,195 @@
+// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::PathBuf;
+
+use clap::Subcommand;
+use clap_num::maybe_hex;
+use vmm::arch::aarch64::regs::Aarch64RegisterVec;
+use vmm::persist::MicrovmState;
+
+use crate::utils::{open_vmstate, save_vmstate, UtilsError};
+
+#[derive(Debug, thiserror::Error)]
+pub enum EditVmStateError {
+    #[error("{0}")]
+    Utils(#[from] UtilsError),
+}
+
+#[derive(Debug, Subcommand)]
+pub enum EditVmStateSubCommand {
+    /// Remove registers from vcpu states.
+    RemoveRegs {
+        /// Set of registers to remove.
+        /// Values should be registers ids as the are defined in KVM.
+        #[arg(value_parser=maybe_hex::<u64>, num_args = 1.., value_delimiter = ' ')]
+        regs: Vec<u64>,
+        /// Path to the vmstate file.
+        #[arg(short, long)]
+        vmstate_path: PathBuf,
+        /// Path of output file.
+        #[arg(short, long)]
+        output_path: PathBuf,
+    },
+}
+
+pub fn edit_vmstate_command(command: EditVmStateSubCommand) -> Result<(), EditVmStateError> {
+    match command {
+        EditVmStateSubCommand::RemoveRegs {
+            regs,
+            vmstate_path,
+            output_path,
+        } => edit(&vmstate_path, &output_path, |state| {
+            remove_regs(state, &regs)
+        })?,
+    }
+    Ok(())
+}
+
+fn edit(
+    vmstate_path: &PathBuf,
+    output_path: &PathBuf,
+    f: impl Fn(MicrovmState) -> Result<MicrovmState, EditVmStateError>,
+) -> Result<(), EditVmStateError> {
+    let (microvm_state, version) = open_vmstate(vmstate_path)?;
+    let microvm_state = f(microvm_state)?;
+    save_vmstate(microvm_state, output_path, version)?;
+    Ok(())
+}
+
+fn remove_regs(
+    mut state: MicrovmState,
+    remove_regs: &[u64],
+) -> Result<MicrovmState, EditVmStateError> {
+    for (i, vcpu_state) in state.vcpu_states.iter_mut().enumerate() {
+        println!("Modifying state for vCPU {i}");
+
+        let mut removed = vec![false; remove_regs.len()];
+        let mut new_regs = Aarch64RegisterVec::default();
+        for reg in vcpu_state.regs.iter().filter(|reg| {
+            if let Some(pos) = remove_regs.iter().position(|r| r == &reg.id) {
+                removed[pos] = true;
+                false
+            } else {
+                true
+            }
+        }) {
+            new_regs.push(reg);
+        }
+        vcpu_state.regs = new_regs;
+        for (reg, removed) in remove_regs.iter().zip(removed.iter()) {
+            print!("Regsiter {reg:#x}: ");
+            match removed {
+                true => println!("removed"),
+                false => println!("not present"),
+            }
+        }
+    }
+    Ok(state)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_remove_regs() {
+        const KVM_REG_SIZE_U8: u64 = 0;
+        const KVM_REG_SIZE_U16: u64 = 0x10000000000000;
+        const KVM_REG_SIZE_U32: u64 = 0x20000000000000;
+
+        use vmm::arch::aarch64::regs::Aarch64RegisterRef;
+        use vmm::vstate::vcpu::aarch64::VcpuState;
+
+        let vcpu_state = VcpuState {
+            regs: {
+                let mut regs = Aarch64RegisterVec::default();
+                let reg_data: u8 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U8,
+                    &reg_data.to_le_bytes(),
+                ));
+                let reg_data: u16 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U16,
+                    &reg_data.to_le_bytes(),
+                ));
+                let reg_data: u32 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U32,
+                    &reg_data.to_le_bytes(),
+                ));
+                regs
+            },
+            ..Default::default()
+        };
+        let state = MicrovmState {
+            vcpu_states: vec![vcpu_state],
+            ..Default::default()
+        };
+
+        let new_state = remove_regs(state, &[KVM_REG_SIZE_U32]).unwrap();
+
+        let expected_vcpu_state = VcpuState {
+            regs: {
+                let mut regs = Aarch64RegisterVec::default();
+                let reg_data: u8 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U8,
+                    &reg_data.to_le_bytes(),
+                ));
+                let reg_data: u16 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U16,
+                    &reg_data.to_le_bytes(),
+                ));
+                regs
+            },
+            ..Default::default()
+        };
+
+        assert_eq!(new_state.vcpu_states[0].regs, expected_vcpu_state.regs);
+    }
+
+    #[test]
+    fn test_remove_non_existed_regs() {
+        const KVM_REG_SIZE_U8: u64 = 0;
+        const KVM_REG_SIZE_U16: u64 = 0x10000000000000;
+        const KVM_REG_SIZE_U32: u64 = 0x20000000000000;
+
+        use vmm::arch::aarch64::regs::Aarch64RegisterRef;
+        use vmm::vstate::vcpu::aarch64::VcpuState;
+
+        let vcpu_state = VcpuState {
+            regs: {
+                let mut regs = Aarch64RegisterVec::default();
+                let reg_data: u8 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U8,
+                    &reg_data.to_le_bytes(),
+                ));
+                let reg_data: u16 = 69;
+                regs.push(Aarch64RegisterRef::new(
+                    KVM_REG_SIZE_U16,
+                    &reg_data.to_le_bytes(),
+                ));
+                regs
+            },
+            ..Default::default()
+        };
+
+        let state_clone = MicrovmState {
+            vcpu_states: vec![vcpu_state.clone()],
+            ..Default::default()
+        };
+
+        let state = MicrovmState {
+            vcpu_states: vec![vcpu_state],
+            ..Default::default()
+        };
+
+        let new_state = remove_regs(state_clone, &[KVM_REG_SIZE_U32]).unwrap();
+
+        assert_eq!(new_state.vcpu_states[0].regs, state.vcpu_states[0].regs);
+    }
+}