From 87ee3351d7894ef5b6b8960dea6487a37d38a870 Mon Sep 17 00:00:00 2001
From: Kalle Wachsmuth <kalle.wachsmuth@gmail.com>
Date: Mon, 1 Apr 2024 15:47:20 +0200
Subject: [PATCH] check block order

---
 src/blob.rs   | 37 +++++++++++++++++++++++++------------
 src/lib.rs    |  6 ++++--
 tests/test.rs | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 60 insertions(+), 15 deletions(-)

diff --git a/src/blob.rs b/src/blob.rs
index b3293f1..b564abe 100644
--- a/src/blob.rs
+++ b/src/blob.rs
@@ -232,23 +232,33 @@ impl Devicetree {
 		}
 
 		let exact_size = self.exact_size() as usize;
-		let (offset, size) = Option::zip(
-			usize::try_from(u32::from_be(self.header().off_dt_struct)).ok(),
+		let (struct_offset, struct_size) = Option::zip(
+			usize::try_from(u32::from_be(self.header().off_dt_struct))
+				.ok()
+				.filter(|&o| o >= Header::SIZE),
 			usize::try_from(u32::from_be(self.header().size_dt_struct)).ok(),
 		)
-		.filter(|&(o, s)| usize::checked_add(o, s).is_some_and(|e| e <= exact_size))
 		.ok_or(BlobError::BlockOutOfBounds)?;
+		let struct_end_offset = usize::checked_add(struct_offset, struct_size)
+			.filter(|&e| e <= exact_size)
+			.ok_or(BlobError::BlockOutOfBounds)?;
 
-		if offset % STRUCT_BLOCK_OPTIMAL_ALIGN != 0 || size % STRUCT_BLOCK_OPTIMAL_ALIGN != 0 {
+		if struct_offset % STRUCT_BLOCK_OPTIMAL_ALIGN != 0
+			|| struct_size % STRUCT_BLOCK_OPTIMAL_ALIGN != 0
+		{
 			return Err(BlobError::UnalignedBlock);
 		}
 
-		if !Option::zip(
-			usize::try_from(u32::from_be(self.header().off_dt_strings)).ok(),
-			usize::try_from(u32::from_be(self.header().size_dt_strings)).ok(),
-		)
-		.and_then(|(o, s)| usize::checked_add(o, s))
-		.is_some_and(|e| e <= exact_size)
+		let strings_offset = usize::try_from(u32::from_be(self.header().off_dt_strings))
+			.map_err(|_| BlobError::BlockOutOfBounds)?;
+		if struct_end_offset > strings_offset {
+			return Err(BlobError::InvalidBlockOrder);
+		}
+
+		if !usize::try_from(u32::from_be(self.header().size_dt_strings))
+			.ok()
+			.and_then(|s| usize::checked_add(strings_offset, s))
+			.is_some_and(|e| e <= exact_size)
 		{
 			return Err(BlobError::BlockOutOfBounds);
 		}
@@ -388,14 +398,17 @@ impl Devicetree {
 			return Err(BlobError::UnalignedBlock);
 		}
 
-		let offset = usize::try_from(offset).map_err(|_| BlobError::BlockOutOfBounds)?;
+		let offset = usize::try_from(offset)
+			.ok()
+			.filter(|&o| o >= Header::SIZE)
+			.ok_or(BlobError::BlockOutOfBounds)?;
 		// type guarantees that `totalsize` is valid and the struct block is in-bounds
 		let end_offset = u32::from_be(self.header().off_dt_struct) as usize;
 		Ok(MemReserveEntries {
 			blob: self
 				.blob
 				.get(offset / DTB_OPTIMAL_ALIGN..end_offset / DTB_OPTIMAL_ALIGN)
-				.ok_or(BlobError::BlockOutOfBounds)?,
+				.ok_or(BlobError::InvalidBlockOrder)?,
 		})
 	}
 }
diff --git a/src/lib.rs b/src/lib.rs
index b8a7427..76be719 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -81,6 +81,7 @@ pub type Result<T, E = Error> = core::result::Result<T, E>;
 pub enum BlobError {
 	BlockOutOfBounds,
 	IncompatibleVersion,
+	InvalidBlockOrder,
 	InvalidPropertyHeader,
 	InvalidRootNode,
 	InvalidString,
@@ -98,6 +99,7 @@ impl Display for BlobError {
 		let description = match *self {
 			Self::BlockOutOfBounds => "block out of bounds",
 			Self::IncompatibleVersion => "incompatible devicetree version",
+			Self::InvalidBlockOrder => "invalid block order",
 			Self::InvalidPropertyHeader => "invalid property header",
 			Self::InvalidRootNode => "invalid root node",
 			Self::InvalidString => "invalid string",
@@ -105,8 +107,8 @@ impl Display for BlobError {
 			Self::NoMagicSignature => "no magic signature",
 			Self::UnalignedBlock => "unaligned block",
 			Self::UnexpectedEnd => "unexpected end",
-			Self::UnexpectedEndToken => "unexpected END token",
-			Self::UnexpectedEndNodeToken => "unexpected END_NODE token",
+			Self::UnexpectedEndToken => "unexpected End token",
+			Self::UnexpectedEndNodeToken => "unexpected EndNode token",
 			Self::UnknownToken => "unknown token",
 		};
 		write!(f, "devicetree blob error: {description}")
diff --git a/tests/test.rs b/tests/test.rs
index 540b003..7c6cad9 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -331,7 +331,7 @@ fn err_invalid_blocks() {
 	assert_matches!(
 		Devicetree::from_slice(&blob! {
 			buf: MINIMAL_BUF,
-			header: { off_mem_rsvmap: 0x8000_0028 },
+			header: { off_mem_rsvmap: 0x20 },
 		})
 		.unwrap()
 		.mem_reserve_entries(),
@@ -353,6 +353,13 @@ fn err_invalid_blocks() {
 		},
 		Some(BlobError::UnalignedBlock),
 	);
+	all_ctors_return(
+		&blob! {
+			buf: MINIMAL_BUF,
+			header: { off_dt_struct: 0x24 },
+		},
+		Some(BlobError::BlockOutOfBounds),
+	);
 	all_ctors_return(
 		&blob! {
 			buf: MINIMAL_BUF,
@@ -399,6 +406,29 @@ fn err_invalid_blocks() {
 	);
 }
 
+#[test]
+fn err_invalid_block_order() {
+	// mem reservation / struct
+	assert_matches!(
+		Devicetree::from_slice(&blob! {
+			buf: MINIMAL_BUF,
+			header: { off_mem_rsvmap: 0x40 },
+		})
+		.unwrap()
+		.mem_reserve_entries(),
+		Err(BlobError::InvalidBlockOrder)
+	);
+
+	// struct / strings
+	all_ctors_return(
+		&blob! {
+			buf: MINIMAL_BUF,
+			header: { off_dt_strings: 0x47 },
+		},
+		Some(BlobError::InvalidBlockOrder),
+	);
+}
+
 #[test]
 fn err_unexpected_end() {
 	// mem reservation