Add tool to predict kernel boot measurements (project-oak#4835)

When we start a new VM we pass the kernel in a bzImage format, but the VMM splits it into 2 parts before passing it to the Stage 0 firmware. So the measurements by stage 0 are of these two derived parts.

This tool applies the same changes to that we can predict what both measurements would be for a specific bzImage kernel.

Cargo.toml

@@ -40,6 +40,7 @@ members = [
   "oak_functions_test_module",
   "oak_functions_test_utils",
   "oak_hello_world_linux_init",
+  "oak_kernel_measurement",
   "oak_launcher_utils",
   "oak_proto_rust",
   "oak_restricted_kernel",
@@ -59,7 +60,7 @@ members = [
   "stage0_dice",
   "testing/oak_echo_service",
   "xtask",
-  "oak_restricted_kernel_sdk_proc_macro"
+  "oak_restricted_kernel_sdk_proc_macro",
 ]
 exclude = [
   "fuzz",

oak_kernel_measurement/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "oak_kernel_measurement"
+version = "0.1.0"
+authors = ["Conrad Grobler <[email protected]>"]
+edition = "2021"
+license = "Apache-2.0"
+
+[dependencies]
+anyhow = "*"
+clap = { version = "*", features = ["derive"] }
+env_logger = "*"
+hex = "*"
+log = "*"
+sha2 = "*"
+zerocopy = "*"

oak_kernel_measurement/README.md

@@ -0,0 +1,18 @@
+# Kernel Measurement Prediction Tool
+
+When booting a VM the kernel and its associated setup data is measured by the
+Stage 0 firmware. The VMM accepts a single bzImage format kernel file. It then
+splits it into two parts, the kernel image and the kernel setup data. The VMM
+also makes modifications to the setup data before passing it to Stage 0 via the
+fw_cfg device.
+
+Stage 0 measures these split, modified components rather than the original
+bzImage kernel. This tool can be used to predict the Stage 0 measurements of
+these components from a bzImage kernel.
+
+The tool can be run using:
+
+```bash
+cargo run --package=oak_kernel_measurement -- \
+    --kernel=oak_containers_kernel/target/bzImage
+```

oak_kernel_measurement/src/main.rs

@@ -0,0 +1,123 @@
+//
+// Copyright 2024 The Project Oak Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+use std::path::PathBuf;
+
+use anyhow::Context;
+use clap::Parser;
+use sha2::{Digest, Sha256};
+use zerocopy::FromBytes;
+
+/// The default workspace-relative path to the Linux Kernel bzImage file.
+const DEFAULT_LINUX_KERNEL: &str = "oak_containers_kernel/target/bzImage";
+
+#[derive(Parser, Clone)]
+#[command(about = "Oak Kernel Measurement Calculator")]
+struct Cli {
+    #[arg(long, help = "The location of the kernel bzImage file")]
+    kernel: Option<PathBuf>,
+}
+
+impl Cli {
+    fn kernel_path(&self) -> PathBuf {
+        self.kernel.clone().unwrap_or_else(|| {
+            format!("{}/{}", env!("WORKSPACE_ROOT"), DEFAULT_LINUX_KERNEL).into()
+        })
+    }
+}
+
+fn main() -> anyhow::Result<()> {
+    env_logger::init();
+    let cli = Cli::parse();
+    let kernel_info = Kernel::load(cli.kernel_path()).context("couldn't load kernel file")?;
+    let mut kernel_hasher = Sha256::new();
+    kernel_hasher.update(&kernel_info.kernel_image);
+    println!(
+        "Kernel Image Measurement: sha2-256:{}",
+        hex::encode(kernel_hasher.finalize())
+    );
+
+    let mut setup_hasher = Sha256::new();
+    setup_hasher.update(&kernel_info.setup_data);
+    println!(
+        "Kernel Setup Data Measurement: sha2-256:{}",
+        hex::encode(setup_hasher.finalize())
+    );
+
+    Ok(())
+}
+
+struct Kernel {
+    setup_data: Vec<u8>,
+    kernel_image: Vec<u8>,
+}
+
+impl Kernel {
+    /// Parses a bzImage kernel file and extracts the kernel image and setup data.
+    ///
+    /// The VMM will parse the the bzImage file and split it into two parts: the setup data (which
+    /// includes the real-mode code) and the 64-bit kernel image. The VMM also updates some fields
+    /// in the setup data header.
+    ///
+    /// See <https://github.com/qemu/qemu/blob/6630bc04bccadcf868165ad6bca5a964bb69b067/hw/i386/x86.c#L795>
+    /// and <https://www.kernel.org/doc/html/v6.7/arch/x86/boot.html>
+    fn from_bz_image(bz_image: &[u8]) -> Self {
+        // The number of 512 byte sectors -1 is stored at offset 0x1F1.
+        let setup_sects = bz_image[0x1F1];
+        // For backwards compatibility, if setup_sects is 0 it should be set to 4.
+        let setup_sects = if setup_sects == 0 { 4 } else { setup_sects };
+        let setup_size = 512 * (setup_sects as usize + 1);
+
+        // The kernel image is just everything after the setup data.
+        let kernel_image = bz_image[setup_size..].to_vec();
+        let mut setup_data = bz_image[..setup_size].to_vec();
+        // The VMM will modify some fields in the setup data header.
+        // The loader type will be set to `QEMU`.
+        setup_data[0x210] = 0xB0;
+        // The load flags will be updated to enable heap usage.
+        setup_data[0x211] |= 0x80;
+        // Set the default command_line location.
+        let default_cmd_line = 0x20000;
+        *u32::mut_from(&mut setup_data[0x228..0x22C]).expect("invalid slice for cmd_line") =
+            default_cmd_line;
+        // Set the offset to the end of the heap.
+        let default_setup = 0x10000;
+        *u16::mut_from(&mut setup_data[0x224..0x226]).expect("invalid slice for heap end ptr") =
+            (default_cmd_line - default_setup - 0x200) as u16;
+
+        // The location and size of the initial RAM disk will depend on the actual initial RAM disk.
+        // To have stable measurements Stage 0 will overwrite these with zeros, before measuring and
+        // then overwrite these with the actual values before booting the kernel.
+        *u32::mut_from(&mut setup_data[0x218..0x21C]).expect("invalid slice for initrd location") =
+            0;
+        *u32::mut_from(&mut setup_data[0x21C..0x220]).expect("invalid slice for initrd size") = 0;
+
+        Self {
+            setup_data,
+            kernel_image,
+        }
+    }
+
+    /// Loads the bzImage-format kernel from the supplied path/
+    fn load(kernel_path: PathBuf) -> anyhow::Result<Self> {
+        let kernel_bytes = std::fs::read(kernel_path).context("couldn't load kernel bzImage")?;
+        log::debug!("Kernel size: {}", kernel_bytes.len());
+        let mut hasher = Sha256::new();
+        hasher.update(&kernel_bytes);
+        log::info!("Kernel digest: sha256:{}", hex::encode(hasher.finalize()));
+        Ok(Self::from_bz_image(&kernel_bytes))
+    }
+}

stage0/src/zero_page.rs

@@ -19,7 +19,7 @@ use core::{ffi::CStr, mem::size_of, slice};
 
 use oak_linux_boot_params::{BootE820Entry, BootParams, E820EntryType};
 use x86_64::PhysAddr;
-use zerocopy::AsBytes;
+use zerocopy::{AsBytes, FromBytes};
 
 use crate::{
     cmos::Cmos,
@@ -75,11 +75,21 @@ impl ZeroPage {
                 .expect("no suitable DMA address available");
             let address = crate::phys_to_virt(start);
 
+            // Safety: we have confirmed that the memory is backed by physical RAM and not currently
+            // used for anything else. We will overwrite all of it, so it does not have to be
+            // initialized.
             let buf = unsafe { slice::from_raw_parts_mut::<u8>(address.as_mut_ptr(), size) };
             let actual_size = fw_cfg
                 .read_file(&file, buf)
                 .expect("could not read setup data");
             assert_eq!(actual_size, size, "setup data did not match expected size");
+            // The initial ram disk location and size are not constant. We will overwrite it later
+            // anyway, so we overwrite it with zeros before measuring so we can get consistent
+            // measurement. See <https://www.kernel.org/doc/html/v6.7/arch/x86/boot.html> for
+            // information on the  field offsets.
+            *u32::mut_from(&mut buf[0x218..0x21C]).expect("invalid slice for initrd location") = 0;
+            *u32::mut_from(&mut buf[0x21C..0x220]).expect("invalid slice for initrd size") = 0;
+
             let measurement = crate::measure_byte_slice(buf);
 
             // The header information starts at offset 0x01F1 from the start of the setup data.