From 771e46c7caa5bf0c040d56bfdcb9d3fd755176bd Mon Sep 17 00:00:00 2001
From: Ivan Velickovic <i.velickovic@unsw.edu.au>
Date: Tue, 3 Oct 2023 20:12:35 +1100
Subject: [PATCH] libvmm: check DTB header when setting up guest images

---
 src/arch/aarch64/linux.c |  9 ++++++++-
 src/dtb.h                | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 src/dtb.h

diff --git a/src/arch/aarch64/linux.c b/src/arch/aarch64/linux.c
index 2f5b6a100..7be8e42f7 100644
--- a/src/arch/aarch64/linux.c
+++ b/src/arch/aarch64/linux.c
@@ -5,6 +5,7 @@
  */
 
 #include "../../util/util.h"
+#include "dtb.h"
 #include "linux.h"
 
 uintptr_t linux_setup_images(uintptr_t ram_start,
@@ -17,7 +18,6 @@ uintptr_t linux_setup_images(uintptr_t ram_start,
                              uintptr_t initrd_dest,
                              size_t initrd_size)
 {
-    // @ivanv: is there a DTB magic to check?
     // @ivanv: is there a initrd magic to check?
     // First we inspect the kernel image header to confirm it is a valid image
     // and to determine where in memory to place the image.
@@ -37,6 +37,13 @@ uintptr_t linux_setup_images(uintptr_t ram_start,
     LOG_VMM("Copying guest kernel image to 0x%x (0x%x bytes)\n", kernel_dest, kernel_size);
     memcpy((char *)kernel_dest, (char *)kernel, kernel_size);
     // Copy the guest device tree blob into the right location
+    // First check that the DTB given is actually a DTB!
+    struct dtb_header *dtb_header = (struct dtb_header *) dtb_src;
+    assert(dtb_check_magic(dtb_header));
+    if (!dtb_check_magic(dtb_header)) {
+        LOG_VMM_ERR("Given DTB does not match DTB magic.\n");
+        return 0;
+    }
     // Linux does not allow the DTB to be greater than 2 megabytes in size.
     assert(dtb_size <= (1 << 21));
     if (dtb_size > (1 << 21)) {
diff --git a/src/dtb.h b/src/dtb.h
new file mode 100644
index 000000000..2f57fcbd8
--- /dev/null
+++ b/src/dtb.h
@@ -0,0 +1,35 @@
+/*
+ * Copyright 2023, UNSW
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <stdint.h>
+
+/*
+ * The following value and header definition are from v0.3, Section 5.2 of the
+ * Devicetree Specification.
+ */
+
+/* Note that this is the little-endian representation, the specification
+ * mentions only the big-endian representation. */
+#define DTB_MAGIC 0xEDFE0DD0
+
+struct dtb_header {
+    uint32_t magic;
+    uint32_t totalsize;
+    uint32_t off_dt_struct;
+    uint32_t off_dt_strings;
+    uint32_t off_mem_rsvmap;
+    uint32_t version;
+    uint32_t last_comp_version;
+    uint32_t boot_cpuid_phys;
+    uint32_t size_dt_strings;
+    uint32_t size_dt_struct;
+};
+
+bool dtb_check_magic(struct dtb_header *h) {
+    return h->magic == DTB_MAGIC;
+}