From e55c413261456c7f1461dbd98f739d4a362bf9e7 Mon Sep 17 00:00:00 2001
From: LoveSy <shana@zju.edu.cn>
Date: Sun, 23 Feb 2025 23:14:28 +0800
Subject: [PATCH] Correctly handle truncated dtb

---
 native/src/boot/dtb.rs | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/native/src/boot/dtb.rs b/native/src/boot/dtb.rs
index f44f0b6ae..38d4750e4 100644
--- a/native/src/boot/dtb.rs
+++ b/native/src/boot/dtb.rs
@@ -1,10 +1,7 @@
 use std::{cell::UnsafeCell, process::exit};
 
 use argh::FromArgs;
-use fdt::{
-    node::{FdtNode, NodeProperty},
-    Fdt,
-};
+use fdt::{node::{FdtNode, NodeProperty}, Fdt, FdtError};
 
 use base::{
     libc::c_char, log_err, map_args, EarlyExitExt, LoggedResult, MappedFile, ResultExt, Utf8CStr,
@@ -171,15 +168,17 @@ fn for_each_fdt<F: FnMut(usize, Fdt) -> LoggedResult<()>>(
         if slice.len() < 40 {
             break;
         }
-        let fdt = Fdt::new(slice)?;
+        let fdt = match Fdt::new(slice) {
+            Err(FdtError::BufferTooSmall) => {
+                eprintln!("dtb.{:04} is truncated", dtb_num);
+                break;
+            },
+            Ok(fdt) => fdt,
+            e => e?,
+        };
 
         let size = fdt.total_size();
 
-        if size > slice.len() {
-            eprintln!("dtb.{:04} is truncated", dtb_num);
-            break;
-        }
-
         f(dtb_num, fdt)?;
 
         dtb_num += 1;