Merge "ComposTestCase: Abort shell command if cd fails"
diff --git a/microdroid/README.md b/microdroid/README.md
index 71be7d0..5e3f586 100644
--- a/microdroid/README.md
+++ b/microdroid/README.md
@@ -8,8 +8,7 @@
 ## Prerequisites
 
 Any 64-bit target (either x86\_64 or arm64) is supported. 32-bit target is not
-supported. Note that we currently don't support user builds; only userdebug
-builds are supported.
+supported.
 
 The only remaining requirement is that `com.android.virt` APEX has to be
 pre-installed. To do this, add the following line in your product makefile.
@@ -18,10 +17,10 @@
 $(call inherit-product, packages/modules/Virtualization/apex/product_packages.mk)
 ```
 
-Build the target after adding the line, and flash it. This step needs to be done
-only once for the target.
+Build the target product after adding the line, and flash it. This step needs
+to be done only once for the target.
 
-If you are using `aosp_oriole` (Pixel 6) or `aosp_cf_x86_64_phone` (Cuttlefish),
+If you are using Pixel 6 and beyond or Cuttlefish (`aosp_cf_x86_64_phone`)
 adding above line is not necessary as it's already done.
 
 ## Building and installing microdroid
@@ -41,8 +40,11 @@
 
 ## Building an app
 
-An app in microdroid is a shared library file embedded in an APK. The shared
-library should have an entry point `AVmPayload_main` as shown below:
+A [vm
+payload](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/vm_payload/)
+is a shared library file that gets executed in microdroid. It is packaged as
+part of an Android application.  The library should have an entry point
+`AVmPayload_main` as shown below:
 
 ```C++
 extern "C" int AVmPayload_main() {
@@ -54,53 +56,22 @@
 
 ```
 cc_library_shared {
-  name: "MyMicrodroidApp",
+  name: "MyMicrodroidPayload",
   srcs: ["**/*.cpp"],
   sdk_version: "current",
 }
 ```
 
-Then you need a configuration file in JSON format that defines what to load and
-execute in microdroid. The name of the file can be anything and you may have
-multiple configuration files if needed.
-
-```json
-{
-  "os": { "name": "microdroid" },
-  "task": {
-    "type": "microdroid_launcher",
-    "command": "MyMicrodroidApp.so"
-  }
-}
-```
-
-The value of `task.command` should match with the name of the shared library
-defined above. If your app requires APEXes to be imported, you can declare the
-list in `apexes` key like following.
-
-```json
-{
-  "os": ...,
-  "task": ...,
-  "apexes": [
-    {"name": "com.android.awesome_apex"}
-  ]
-}
-```
-
-Embed the shared library and the VM configuration file in an APK:
+Embed the shared library file in an APK:
 
 ```
 android_app {
   name: "MyApp",
-  srcs: ["**/*.java"], // if there is any java code
-  jni_libs: ["MyMicrodroidApp"],
+  srcs: ["**/*.java"],
+  jni_libs: ["MyMicrodroidPayload"],
   use_embedded_native_libs: true,
   sdk_version: "current",
 }
-
-// The VM configuration file can be embedded by simply placing it at `./assets`
-// directory.
 ```
 
 Finally, you build the APK.
@@ -109,7 +80,7 @@
 TARGET_BUILD_APPS=MyApp m apps_only dist
 ```
 
-## Running the app on microdroid
+## Running the VM payload on microdroid
 
 First of all, install the APK to the target device.
 
@@ -117,22 +88,16 @@
 adb install out/dist/MyApp.apk
 ```
 
-`ALL_CAP`s below are placeholders. They need to be replaced with correct
-values:
+There are two ways start a VM and run the payload in it.
 
-* `VM_CONFIG_FILE`: the name of the VM config file that you embedded in the APK.
-  (e.g. `vm_config.json`)
-* `PACKAGE_NAME_OF_YOUR_APP`: package name of your app (e.g. `com.acme.app`).
-* `PATH_TO_YOUR_APP`: path to the installed APK on the device. Can be obtained
-  via the following command.
-  ```sh
-  adb shell pm path PACKAGE_NAME_OF_YOUR_APP
-  ```
-  It shall report a cryptic path similar to `/data/app/~~OgZq==/com.acme.app-HudMahQ==/base.apk`.
+* By manually invoking the `vm` tool via `adb shell`.
+* Calling APIs programmatically in the Java app.
+
+### Using `vm` tool
 
 Execute the following commands to launch a VM. The VM will boot to microdroid
-and then automatically execute your app (the shared library
-`MyMicrodroidApp.so`).
+and then automatically execute your payload (the shared library
+`MyMicrodroidPayload.so`).
 
 ```sh
 TEST_ROOT=/data/local/tmp/virt
@@ -142,23 +107,36 @@
 PATH_TO_YOUR_APP \
 $TEST_ROOT/MyApp.apk.idsig \
 $TEST_ROOT/instance.img \
---config-path assets/VM_CONFIG_FILE
+--payload-binary-name MyMicrodroidPayload.so
 ```
 
-The last command lets you know the CID assigned to the VM. The console output
-from the VM is stored to `$TEST_ROOT/console.txt` and logcat is stored to
-`$TEST_ROOT/log.txt` file for debugging purpose. If you omit `--log` or
-`--console` option, they will be emitted to the current console.
+`ALL_CAP`s below are placeholders. They need to be replaced with correct
+values:
 
-Stopping the VM can be done as follows:
+* `PACKAGE_NAME_OF_YOUR_APP`: package name of your app (e.g. `com.acme.app`).
+* `PATH_TO_YOUR_APP`: path to the installed APK on the device. Can be obtained
+  via the following command.
+  ```sh
+  adb shell pm path PACKAGE_NAME_OF_YOUR_APP
+  ```
+  It shall report a cryptic path similar to `/data/app/~~OgZq==/com.acme.app-HudMahQ==/base.apk`.
 
-```sh
-adb shell /apex/com.android.virt/bin/vm stop $CID
-```
+The console output from the VM is stored to `$TEST_ROOT/console.txt` and logcat
+is stored to `$TEST_ROOT/log.txt` file for debugging purpose. If you omit
+`--log` or `--console` option, the console output will be emitted to the
+current console and the logcat logs are sent to the main logcat in Android.
 
-, where `$CID` is the reported CID value. This works only when the `vm` was
-invoked with the `--daemonize` flag. If the flag was not used, press Ctrl+C on
-the console where the `vm run-app` command was invoked.
+Stopping the VM can be done by pressing `Ctrl+C`.
+
+### Using the APIs
+
+Use the [Android Virtualization Framework Java
+APIs](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/javalib/api/system-current.txt)
+in your app to create a microdroid VM and run payload in it. The APIs currently
+are @SystemApi, thus available only to privileged apps.
+
+If you are looking for an example usage of the APIs, you may refer to the [demo
+app](https://android.googlesource.com/platform/packages/modules/Virtualization/+/refs/heads/master/demo/).
 
 ## Debuggable microdroid
 
diff --git a/microdroid/init.rc b/microdroid/init.rc
index 29f8970..871db94 100644
--- a/microdroid/init.rc
+++ b/microdroid/init.rc
@@ -146,6 +146,7 @@
     capabilities CHOWN DAC_OVERRIDE DAC_READ_SEARCH FOWNER SYS_ADMIN
 
 service ueventd /system/bin/ueventd
+    user root
     class core
     critical
     seclabel u:r:ueventd:s0
@@ -161,6 +162,7 @@
     setenv HOSTNAME console
 
 service init_debug_policy /system/bin/init_debug_policy
+    user root
     oneshot
     disabled
     stdio_to_kmsg
diff --git a/microdroid/payload/mk_payload.cc b/microdroid/payload/mk_payload.cc
deleted file mode 100644
index d31333f..0000000
--- a/microdroid/payload/mk_payload.cc
+++ /dev/null
@@ -1,305 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * 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.
- */
-
-#include <sys/stat.h>
-#include <sys/types.h>
-#include <unistd.h>
-
-#include <fstream>
-#include <iostream>
-#include <optional>
-#include <string>
-#include <vector>
-
-#include <android-base/file.h>
-#include <android-base/result.h>
-#include <image_aggregator.h>
-#include <json/json.h>
-
-#include "microdroid/metadata.h"
-
-using android::base::Dirname;
-using android::base::ErrnoError;
-using android::base::Error;
-using android::base::Result;
-using android::base::unique_fd;
-using android::microdroid::ApexPayload;
-using android::microdroid::ApkPayload;
-using android::microdroid::Metadata;
-using android::microdroid::WriteMetadata;
-
-using cuttlefish::AlignToPartitionSize;
-using cuttlefish::CreateCompositeDisk;
-using cuttlefish::kLinuxFilesystem;
-using cuttlefish::MultipleImagePartition;
-
-Result<uint32_t> GetFileSize(const std::string& path) {
-    struct stat st;
-    if (lstat(path.c_str(), &st) == -1) {
-        return ErrnoError() << "Can't lstat " << path;
-    }
-    return static_cast<uint32_t>(st.st_size);
-}
-
-std::string RelativeTo(const std::string& path, const std::string& dirname) {
-    bool is_absolute = !path.empty() && path[0] == '/';
-    if (is_absolute || dirname == ".") {
-        return path;
-    } else {
-        return dirname + "/" + path;
-    }
-}
-
-// Returns `append` is appended to the end of filename preserving the extension.
-std::string AppendFileName(const std::string& filename, const std::string& append) {
-    size_t pos = filename.find_last_of('.');
-    if (pos == std::string::npos) {
-        return filename + append;
-    } else {
-        return filename.substr(0, pos) + append + filename.substr(pos);
-    }
-}
-
-struct ApexConfig {
-    std::string name; // the apex name
-    std::string path; // the path to the apex file
-                      // absolute or relative to the config file
-};
-
-struct ApkConfig {
-    std::string name;
-    std::string path;
-    std::string idsig_path;
-};
-
-struct Config {
-    std::string dirname; // config file's direname to resolve relative paths in the config
-
-    std::vector<ApexConfig> apexes;
-    std::optional<ApkConfig> apk;
-    // This is a path in the guest side
-    std::optional<std::string> payload_config_path;
-};
-
-#define DO(expr) \
-    if (auto res = (expr); !res.ok()) return res.error()
-
-Result<void> ParseJson(const Json::Value& value, std::string& s) {
-    if (!value.isString()) {
-        return Error() << "should be a string: " << value;
-    }
-    s = value.asString();
-    return {};
-}
-
-template <typename T>
-Result<void> ParseJson(const Json::Value& value, std::optional<T>& s) {
-    if (value.isNull()) {
-        s.reset();
-        return {};
-    }
-    s.emplace();
-    return ParseJson(value, *s);
-}
-
-template <typename T>
-Result<void> ParseJson(const Json::Value& values, std::vector<T>& parsed) {
-    for (const Json::Value& value : values) {
-        T t;
-        DO(ParseJson(value, t));
-        parsed.push_back(std::move(t));
-    }
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, ApexConfig& apex_config) {
-    DO(ParseJson(value["name"], apex_config.name));
-    DO(ParseJson(value["path"], apex_config.path));
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, ApkConfig& apk_config) {
-    DO(ParseJson(value["name"], apk_config.name));
-    DO(ParseJson(value["path"], apk_config.path));
-    DO(ParseJson(value["idsig_path"], apk_config.idsig_path));
-    return {};
-}
-
-Result<void> ParseJson(const Json::Value& value, Config& config) {
-    DO(ParseJson(value["apexes"], config.apexes));
-    DO(ParseJson(value["apk"], config.apk));
-    DO(ParseJson(value["payload_config_path"], config.payload_config_path));
-    return {};
-}
-
-Result<Config> LoadConfig(const std::string& config_file) {
-    std::ifstream in(config_file);
-    Json::CharReaderBuilder builder;
-    Json::Value root;
-    Json::String errs;
-    if (!parseFromStream(builder, in, &root, &errs)) {
-        return Error() << errs;
-    }
-
-    Config config;
-    config.dirname = Dirname(config_file);
-    DO(ParseJson(root, config));
-    return config;
-}
-
-#undef DO
-
-Result<void> MakeMetadata(const Config& config, const std::string& filename) {
-    Metadata metadata;
-    metadata.set_version(1);
-
-    for (const auto& apex_config : config.apexes) {
-        auto* apex = metadata.add_apexes();
-        apex->set_name(apex_config.name);
-        apex->set_partition_name(apex_config.name);
-        apex->set_is_factory(true);
-    }
-
-    if (config.apk.has_value()) {
-        auto* apk = metadata.mutable_apk();
-        apk->set_name(config.apk->name);
-        apk->set_payload_partition_name("microdroid-apk");
-        apk->set_idsig_partition_name("microdroid-apk-idsig");
-    }
-
-    if (config.payload_config_path.has_value()) {
-        *metadata.mutable_config_path() = config.payload_config_path.value();
-    }
-
-    std::ofstream out(filename);
-    return WriteMetadata(metadata, out);
-}
-
-// fill zeros to align |file_path|'s size to BLOCK_SIZE(4096) boundary.
-// return true when the filler is needed.
-Result<bool> ZeroFiller(const std::string& file_path, const std::string& filler_path) {
-    auto file_size = GetFileSize(file_path);
-    if (!file_size.ok()) {
-        return file_size.error();
-    }
-    auto disk_size = AlignToPartitionSize(*file_size);
-    if (disk_size <= *file_size) {
-        return false;
-    }
-    unique_fd fd(TEMP_FAILURE_RETRY(open(filler_path.c_str(), O_CREAT | O_WRONLY | O_TRUNC, 0600)));
-    if (fd.get() == -1) {
-        return ErrnoError() << "open(" << filler_path << ") failed.";
-    }
-    if (ftruncate(fd.get(), disk_size - *file_size) == -1) {
-        return ErrnoError() << "ftruncate(" << filler_path << ") failed.";
-    }
-    return true;
-}
-
-Result<void> MakePayload(const Config& config, const std::string& metadata_file,
-                         const std::string& output_file) {
-    std::vector<MultipleImagePartition> partitions;
-
-    int filler_count = 0;
-    auto add_partition = [&](auto partition_name, auto file_path) -> Result<void> {
-        std::vector<std::string> image_files{file_path};
-
-        std::string filler_path =
-                AppendFileName(output_file, "-filler-" + std::to_string(filler_count++));
-        if (auto ret = ZeroFiller(file_path, filler_path); !ret.ok()) {
-            return ret.error();
-        } else if (*ret) {
-            image_files.push_back(filler_path);
-        }
-        partitions.push_back(MultipleImagePartition{
-                .label = partition_name,
-                .image_file_paths = image_files,
-                .type = kLinuxFilesystem,
-                .read_only = true,
-        });
-        return {};
-    };
-
-    // put metadata at the first partition
-    partitions.push_back(MultipleImagePartition{
-            .label = "payload-metadata",
-            .image_file_paths = {metadata_file},
-            .type = kLinuxFilesystem,
-            .read_only = true,
-    });
-    // put apexes at the subsequent partitions
-    for (size_t i = 0; i < config.apexes.size(); i++) {
-        const auto& apex_config = config.apexes[i];
-        std::string apex_path = RelativeTo(apex_config.path, config.dirname);
-        if (auto ret = add_partition("microdroid-apex-" + std::to_string(i), apex_path);
-            !ret.ok()) {
-            return ret.error();
-        }
-    }
-    // put apk and its idsig
-    if (config.apk.has_value()) {
-        std::string apk_path = RelativeTo(config.apk->path, config.dirname);
-        if (auto ret = add_partition("microdroid-apk", apk_path); !ret.ok()) {
-            return ret.error();
-        }
-        std::string idsig_path = RelativeTo(config.apk->idsig_path, config.dirname);
-        if (auto ret = add_partition("microdroid-apk-idsig", idsig_path); !ret.ok()) {
-            return ret.error();
-        }
-    }
-
-    const std::string gpt_header = AppendFileName(output_file, "-header");
-    const std::string gpt_footer = AppendFileName(output_file, "-footer");
-    CreateCompositeDisk(partitions, gpt_header, gpt_footer, output_file);
-    return {};
-}
-
-int main(int argc, char** argv) {
-    if (argc < 3 || argc > 4) {
-        std::cerr << "Usage: " << argv[0] << " [--metadata-only] <config> <output>\n";
-        return 1;
-    }
-    int arg_index = 1;
-    bool metadata_only = false;
-    if (strcmp(argv[arg_index], "--metadata-only") == 0) {
-        metadata_only = true;
-        arg_index++;
-    }
-
-    auto config = LoadConfig(argv[arg_index++]);
-    if (!config.ok()) {
-        std::cerr << "bad config: " << config.error() << '\n';
-        return 1;
-    }
-
-    const std::string output_file(argv[arg_index++]);
-    const std::string metadata_file =
-            metadata_only ? output_file : AppendFileName(output_file, "-metadata");
-
-    if (const auto res = MakeMetadata(*config, metadata_file); !res.ok()) {
-        std::cerr << res.error() << '\n';
-        return 1;
-    }
-    if (metadata_only) {
-        return 0;
-    }
-    if (const auto res = MakePayload(*config, metadata_file, output_file); !res.ok()) {
-        std::cerr << res.error() << '\n';
-        return 1;
-    }
-
-    return 0;
-}
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 7ea1189..0ae2203 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -17,6 +17,8 @@
         "libaarch64_paging",
         "libbssl_ffi_nostd",
         "libbuddy_system_allocator",
+        "libciborium_nostd",
+        "libciborium_io_nostd",
         "libdiced_open_dice_nostd",
         "libfdtpci",
         "libhyp",
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index a773f1a..1c22861 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -37,6 +37,7 @@
 mod virtio;
 
 use alloc::boxed::Box;
+use alloc::string::ToString;
 
 use crate::dice::PartialInputs;
 use crate::entry::RebootReason;
@@ -46,6 +47,7 @@
 use crate::instance::get_or_generate_instance_salt;
 use crate::memory::MemoryTracker;
 use crate::virtio::pci;
+use ciborium::{de::from_reader, value::Value};
 use diced_open_dice::bcc_handover_main_flow;
 use diced_open_dice::bcc_handover_parse;
 use diced_open_dice::DiceArtifacts;
@@ -58,6 +60,19 @@
 
 const NEXT_BCC_SIZE: usize = GUEST_PAGE_SIZE;
 
+type CiboriumError = ciborium::de::Error<ciborium_io::EndOfFile>;
+
+/// Decodes the provided binary CBOR-encoded value and returns a
+/// ciborium::Value struct wrapped in Result.
+fn value_from_bytes(mut bytes: &[u8]) -> Result<Value, CiboriumError> {
+    let value = from_reader(&mut bytes)?;
+    // Ciborium tries to read one Value, but doesn't care if there is trailing data. We do.
+    if !bytes.is_empty() {
+        return Err(CiboriumError::Semantic(Some(0), "unexpected trailing data".to_string()));
+    }
+    Ok(value)
+}
+
 fn main(
     fdt: &mut Fdt,
     signed_kernel: &[u8],
@@ -81,6 +96,18 @@
     })?;
     trace!("BCC: {bcc_handover:x?}");
 
+    // Minimal BCC verification - check the BCC exists & is valid CBOR.
+    // TODO(alanstokes): Do something more useful.
+    if let Some(bytes) = bcc_handover.bcc() {
+        let _ = value_from_bytes(bytes).map_err(|e| {
+            error!("Invalid BCC: {e:?}");
+            RebootReason::InvalidBcc
+        })?;
+    } else {
+        error!("Missing BCC");
+        return Err(RebootReason::InvalidBcc);
+    }
+
     // Set up PCI bus for VirtIO devices.
     let pci_info = PciInfo::from_fdt(fdt).map_err(handle_pci_error)?;
     debug!("PCI: {:#x?}", pci_info);
diff --git a/pvmfw/src/virtio/hal.rs b/pvmfw/src/virtio/hal.rs
index 5f70b33..7598a55 100644
--- a/pvmfw/src/virtio/hal.rs
+++ b/pvmfw/src/virtio/hal.rs
@@ -9,7 +9,21 @@
 
 pub struct HalImpl;
 
-impl Hal for HalImpl {
+/// Implements the `Hal` trait for `HalImpl`.
+///
+/// # Safety
+///
+/// Callers of this implementatation must follow the safety requirements documented for the unsafe
+/// methods.
+unsafe impl Hal for HalImpl {
+    /// Allocates the given number of contiguous physical pages of DMA memory for VirtIO use.
+    ///
+    /// # Implementation Safety
+    ///
+    /// `dma_alloc` ensures the returned DMA buffer is not aliased with any other allocation or
+    ///  reference in the program until it is deallocated by `dma_dealloc` by allocating a unique
+    ///  block of memory using `alloc_shared` and returning a non-null pointer to it that is
+    ///  aligned to `PAGE_SIZE`.
     fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let size = pages * PAGE_SIZE;
@@ -19,7 +33,7 @@
         (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
+    unsafe fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
         let size = pages * PAGE_SIZE;
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
@@ -30,7 +44,13 @@
         0
     }
 
-    fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
+    /// Converts a physical address used for MMIO to a virtual address which the driver can access.
+    ///
+    /// # Implementation Safety
+    ///
+    /// `mmio_phys_to_virt` satisfies the requirement by checking that the mapped memory region
+    /// is within the PCI MMIO range.
+    unsafe fn mmio_phys_to_virt(paddr: PhysAddr, size: usize) -> NonNull<u8> {
         let pci_info = PCI_INFO.get().expect("VirtIO HAL used before PCI_INFO was initialised");
         // Check that the region is within the PCI MMIO range that we read from the device tree. If
         // not, the host is probably trying to do something malicious.
@@ -48,7 +68,7 @@
         phys_to_virt(paddr)
     }
 
-    fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
+    unsafe fn share(buffer: NonNull<[u8]>, direction: BufferDirection) -> PhysAddr {
         let size = buffer.len();
 
         // TODO: Copy to a pre-shared region rather than allocating and sharing each time.
@@ -63,7 +83,7 @@
         virt_to_phys(copy)
     }
 
-    fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
+    unsafe fn unshare(paddr: PhysAddr, buffer: NonNull<[u8]>, direction: BufferDirection) {
         let vaddr = phys_to_virt(paddr);
         let size = buffer.len();
         if direction == BufferDirection::DeviceToDriver {
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 99e07b6..3398d50 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -28,9 +28,10 @@
     paging::{Attributes, MemoryRegion},
 };
 use buddy_system_allocator::LockedHeap;
+use core::ops::Range;
 use hyp::get_hypervisor;
 use log::{debug, error, info};
-use vmbase::{main, power::reboot};
+use vmbase::{layout, main, power::reboot};
 
 const SZ_1K: usize = 1024;
 const SZ_64K: usize = 64 * SZ_1K;
@@ -63,16 +64,8 @@
 
 static mut HEAP: [u8; SZ_64K] = [0; SZ_64K];
 
-unsafe fn kimg_ptr(sym: &u8) -> *const u8 {
-    sym as *const u8
-}
-
-unsafe fn kimg_addr(sym: &u8) -> usize {
-    kimg_ptr(sym) as usize
-}
-
-unsafe fn kimg_region(begin: &u8, end: &u8) -> MemoryRegion {
-    MemoryRegion::new(kimg_addr(begin), kimg_addr(end))
+fn into_memreg(r: &Range<usize>) -> MemoryRegion {
+    MemoryRegion::new(r.start, r.end)
 }
 
 fn init_heap() {
@@ -86,11 +79,9 @@
 fn init_kernel_pgt(pgt: &mut IdMap) -> Result<()> {
     // The first 1 GiB of address space is used by crosvm for MMIO.
     let reg_dev = MemoryRegion::new(0, SZ_1G);
-    // SAFETY: Taking addresses of kernel image sections to set up page table
-    // mappings. Not taking ownerhip of the memory.
-    let reg_text = unsafe { kimg_region(&text_begin, &text_end) };
-    let reg_rodata = unsafe { kimg_region(&rodata_begin, &rodata_end) };
-    let reg_data = unsafe { kimg_region(&data_begin, &boot_stack_end) };
+    let reg_text = into_memreg(&layout::text_range());
+    let reg_rodata = into_memreg(&layout::rodata_range());
+    let reg_data = into_memreg(&layout::writable_region());
 
     debug!("Preparing kernel page table.");
     debug!("  dev:    {}-{}", reg_dev.start(), reg_dev.end());
@@ -143,13 +134,4 @@
     }
 }
 
-extern "C" {
-    static text_begin: u8;
-    static text_end: u8;
-    static rodata_begin: u8;
-    static rodata_end: u8;
-    static data_begin: u8;
-    static boot_stack_end: u8;
-}
-
 main!(main);
diff --git a/vmbase/example/src/pci.rs b/vmbase/example/src/pci.rs
index 117cbc8..41a3ff4 100644
--- a/vmbase/example/src/pci.rs
+++ b/vmbase/example/src/pci.rs
@@ -98,7 +98,7 @@
 
 struct HalImpl;
 
-impl Hal for HalImpl {
+unsafe impl Hal for HalImpl {
     fn dma_alloc(pages: usize, _direction: BufferDirection) -> (PhysAddr, NonNull<u8>) {
         debug!("dma_alloc: pages={}", pages);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
@@ -110,7 +110,7 @@
         (paddr, vaddr)
     }
 
-    fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
+    unsafe fn dma_dealloc(paddr: PhysAddr, vaddr: NonNull<u8>, pages: usize) -> i32 {
         debug!("dma_dealloc: paddr={:#x}, pages={}", paddr, pages);
         let layout = Layout::from_size_align(pages * PAGE_SIZE, PAGE_SIZE).unwrap();
         // Safe because the memory was allocated by `dma_alloc` above using the same allocator, and
@@ -121,17 +121,17 @@
         0
     }
 
-    fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
+    unsafe fn mmio_phys_to_virt(paddr: PhysAddr, _size: usize) -> NonNull<u8> {
         NonNull::new(paddr as _).unwrap()
     }
 
-    fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
+    unsafe fn share(buffer: NonNull<[u8]>, _direction: BufferDirection) -> PhysAddr {
         let vaddr = buffer.cast();
         // Nothing to do, as the host already has access to all memory.
         virt_to_phys(vaddr)
     }
 
-    fn unshare(_paddr: PhysAddr, _buffer: NonNull<[u8]>, _direction: BufferDirection) {
+    unsafe fn unshare(_paddr: PhysAddr, _buffer: NonNull<[u8]>, _direction: BufferDirection) {
         // Nothing to do, as the host already has access to all memory and we didn't copy the buffer
         // anywhere else.
     }
diff --git a/vmbase/src/layout.rs b/vmbase/src/layout.rs
index b0a5173..b5ab449 100644
--- a/vmbase/src/layout.rs
+++ b/vmbase/src/layout.rs
@@ -14,53 +14,69 @@
 
 //! Memory layout.
 
-use crate::linker;
 use core::ops::Range;
 use core::ptr::addr_of;
 
+/// Get an address from a linker-defined symbol.
+#[macro_export]
+macro_rules! linker_addr {
+    ($symbol:ident) => {{
+        unsafe { addr_of!($crate::linker::$symbol) as usize }
+    }};
+}
+
+/// Get the address range between a pair of linker-defined symbols.
+#[macro_export]
+macro_rules! linker_region {
+    ($begin:ident,$end:ident) => {{
+        let start = linker_addr!($begin);
+        let end = linker_addr!($end);
+
+        start..end
+    }};
+}
+
 /// Memory reserved for the DTB.
 pub fn dtb_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::dtb_begin) as usize)..(addr_of!(linker::dtb_end) as usize) }
+    linker_region!(dtb_begin, dtb_end)
 }
 
 /// Executable code.
 pub fn text_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::text_begin) as usize)..(addr_of!(linker::text_end) as usize) }
+    linker_region!(text_begin, text_end)
 }
 
 /// Read-only data.
 pub fn rodata_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::rodata_begin) as usize)..(addr_of!(linker::rodata_end) as usize) }
+    linker_region!(rodata_begin, rodata_end)
 }
 
 /// Initialised writable data.
 pub fn data_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::data_begin) as usize)..(addr_of!(linker::data_end) as usize) }
+    linker_region!(data_begin, data_end)
 }
 
 /// Zero-initialised writable data.
 pub fn bss_range() -> Range<usize> {
-    unsafe { (addr_of!(linker::bss_begin) as usize)..(addr_of!(linker::bss_end) as usize) }
+    linker_region!(bss_begin, bss_end)
 }
 
 /// Writable data region for the stack.
 pub fn boot_stack_range() -> Range<usize> {
-    unsafe {
-        (addr_of!(linker::boot_stack_begin) as usize)..(addr_of!(linker::boot_stack_end) as usize)
-    }
+    linker_region!(boot_stack_begin, boot_stack_end)
 }
 
 /// Writable data, including the stack.
 pub fn writable_region() -> Range<usize> {
-    data_range().start..boot_stack_range().end
+    linker_region!(data_begin, boot_stack_end)
 }
 
 /// Read-write data (original).
 pub fn data_load_address() -> usize {
-    unsafe { addr_of!(linker::data_lma) as usize }
+    linker_addr!(data_lma)
 }
 
 /// End of the binary image.
 pub fn binary_end() -> usize {
-    unsafe { addr_of!(linker::bin_end) as usize }
+    linker_addr!(bin_end)
 }