metropolis: implement and use A/B preloader

This switches over from using the EFI built-in bootloader for A/B
updates to using our own EFI preloader due to significant issues with
in-the-wild EFI implementations.  It is a very minimal design relying
on a single Protobuf state file instead of EFI variables.

Change-Id: Ieebd0a8172ebe3f44c69b3e8c278c53d3fe2eeb4
Reviewed-on: https://review.monogon.dev/c/monogon/+/2203
Tested-by: Jenkins CI
Reviewed-by: Serge Bazanski <serge@monogon.tech>
diff --git a/metropolis/node/core/abloader/main.rs b/metropolis/node/core/abloader/main.rs
new file mode 100644
index 0000000..0068f95
--- /dev/null
+++ b/metropolis/node/core/abloader/main.rs
@@ -0,0 +1,170 @@
+#![no_main]
+#![no_std]
+
+extern crate alloc;
+
+use alloc::vec::Vec;
+use core::result::Result;
+use core::fmt;
+use prost::Message;
+use uefi::fs::FileSystem;
+use uefi::proto::device_path::build::media::FilePath;
+use uefi::proto::device_path::build::DevicePathBuilder;
+use uefi::proto::device_path::{DeviceSubType, DeviceType, LoadedImageDevicePath};
+use uefi::table::boot;
+use uefi::{prelude::*, CStr16};
+use uefi_services::println;
+
+use abloader_proto::monogon::metropolis::node::core::abloader;
+
+const A_LOADER_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\boot-a.efi");
+const B_LOADER_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\boot-b.efi");
+
+const LOADER_STATE_PATH: &CStr16 = cstr16!("\\EFI\\metropolis\\loader_state.pb");
+
+enum ValidSlot {
+    A,
+    B,
+}
+
+impl ValidSlot {
+    // other returns B if the value is A and A if the value is B.
+    fn other(&self) -> Self {
+        match self {
+            ValidSlot::A => ValidSlot::B,
+            ValidSlot::B => ValidSlot::A,
+        }
+    }
+    // path returns the path to the slot's EFI payload.
+    fn path(&self) -> &'static CStr16 {
+        match self {
+            ValidSlot::A => A_LOADER_PATH,
+            ValidSlot::B => B_LOADER_PATH,
+        }
+    }
+}
+
+enum ReadLoaderStateError {
+    FSReadError(uefi::fs::Error),
+    DecodeError(prost::DecodeError),
+}
+
+impl fmt::Display for ReadLoaderStateError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+           ReadLoaderStateError::FSReadError(e) => write!(f, "while reading state file: {}", e),
+           ReadLoaderStateError::DecodeError(e) => write!(f, "while decoding state file contents: {}", e),
+        }
+    }
+}
+
+fn read_loader_state(fs: &mut FileSystem) -> Result<abloader::AbLoaderData, ReadLoaderStateError> {
+    let state_raw = fs.read(&LOADER_STATE_PATH).map_err(|e| ReadLoaderStateError::FSReadError(e))?;
+    abloader::AbLoaderData::decode(state_raw.as_slice()).map_err(|e| ReadLoaderStateError::DecodeError(e))
+}
+
+fn load_slot_image(slot: &ValidSlot, boot_services: &BootServices) -> uefi::Result<Handle> {
+    let mut storage = Vec::new();
+
+    // Build the path to the slot payload. This takes the path to the loader
+    // itself, strips off the file path and following element(s) and appends
+    // the path to the correct slot payload.
+    let new_image_path = {
+        let loaded_image_device_path = boot_services
+            .open_protocol_exclusive::<LoadedImageDevicePath>(boot_services.image_handle())?;
+
+        let mut builder = DevicePathBuilder::with_vec(&mut storage);
+
+        for node in loaded_image_device_path.node_iter() {
+            if node.full_type() == (DeviceType::MEDIA, DeviceSubType::MEDIA_FILE_PATH) {
+                break;
+            }
+
+            builder = builder.push(&node).unwrap();
+        }
+
+        builder = builder
+            .push(&FilePath {
+                path_name: slot.path(),
+            })
+            .unwrap();
+
+        builder.finalize().unwrap()
+    };
+
+    boot_services
+        .load_image(
+            boot_services.image_handle(),
+            boot::LoadImageSource::FromDevicePath {
+                device_path: new_image_path,
+                from_boot_manager: false,
+            },
+        )
+}
+
+#[entry]
+fn main(_handle: Handle, mut system_table: SystemTable<Boot>) -> Status {
+    uefi_services::init(&mut system_table).unwrap();
+
+    let boot_services = system_table.boot_services();
+
+    let boot_slot_raw = {
+        let mut esp_fs = boot_services
+            .get_image_file_system(boot_services.image_handle())
+            .expect("image filesystem not available");
+
+        let mut loader_data = match read_loader_state(&mut esp_fs) {
+            Ok(d) => d, 
+            Err(e) => {
+                println!("Unable to load A/B loader state, using default slot A: {}", e);
+                abloader::AbLoaderData {
+                    active_slot: abloader::Slot::A.into(),
+                    next_slot: abloader::Slot::None.into(),
+                }
+            }
+        };
+
+        // If next_slot is set, use it as slot to boot but clear it in the
+        // state file as the next boot should not use it again. If it should
+        // be permanently activated, it is the OS's job to put it into 
+        if loader_data.next_slot != abloader::Slot::None.into() {
+            let next_slot = loader_data.next_slot;
+            loader_data.next_slot = abloader::Slot::None.into();
+            let new_loader_data = loader_data.encode_to_vec();
+            esp_fs
+                .write(&LOADER_STATE_PATH, new_loader_data)
+                .expect("failed to write back abdata");
+            next_slot
+        } else {
+            loader_data.active_slot
+        }
+    };
+
+    let boot_slot = match abloader::Slot::try_from(boot_slot_raw) {
+        Ok(abloader::Slot::A) => ValidSlot::A,
+        Ok(abloader::Slot::B) => ValidSlot::B,
+        _ => {
+            println!("Invalid slot ({}) active, falling back to A", boot_slot_raw);
+            ValidSlot::A
+        }
+    };
+
+    let payload_image = match load_slot_image(&boot_slot, boot_services) {
+        Ok(img) => img,
+        Err(e) => {
+            println!("Error loading intended slot, falling back to other slot: {}", e);
+            match load_slot_image(&boot_slot.other(), boot_services) {
+                Ok(img) => img,
+                Err(e) => {
+                    panic!("Loading from both slots failed, second slot error: {}", e);
+                },
+            }
+        }
+    };
+
+    // Boot the payload.
+    boot_services
+        .start_image(payload_image)
+        .expect("failed to start payload");
+    Status::SUCCESS
+}