libufdt: extract_dtb: Check tree is valid before extraction.

Ignore-AOSP-First: Security
Bug: 261442369
Test: mmma system/libufdt
Test: extract_dtb with corrupt tree in image
Test: extract_dtb with valid trees in image
Test: system/libufdt/tests/run_tests.sh
Change-Id: I3d2b29e815323c1d8d3d0d2087ce70d04ab35df4
(cherry picked from commit f2e86523500a4346065e109223dc2c6124188413)

1 files changed, 33 insertions, 20 deletions

diff --git a/tests/src/extract_dtb.c b/tests/src/extract_dtb.c
index 8d74324..30cd1b1 100644
--- a/tests/src/extract_dtb.c
+++ b/tests/src/extract_dtb.c
@@ -22,8 +22,7 @@
 
 #include "util.h"
 
-
-int find_dtb_header_pos(const char *buf, size_t buf_size) {
+static int find_dtb_header_pos(const char *buf, size_t buf_size) {
   if (buf == NULL || buf_size == 0) {
     return -1;
   }
@@ -44,34 +43,44 @@ int find_dtb_header_pos(const char *buf, size_t buf_size) {
   return pos;
 }
 
-int find_and_write_dtb(const char *filename,
-                       const char *buf, size_t buf_size) {
+static int find_and_write_dtb(const char *filename, const char *buf,
+                              size_t buf_size) {
   int tag_pos = find_dtb_header_pos(buf, buf_size);
   if (tag_pos < 0) {
-    goto end;
+    return -1;
   }
 
+  buf_size -= tag_pos;
+
+  // Allocate and copy into new buffer to fix memory alignment
+  char *fdt_ptr = malloc(buf_size);
+  if (!fdt_ptr) {
+    fprintf(stderr, "malloc(%u) failed.\n", buf_size - tag_pos);
+    goto error;
+  }
+
+  memcpy(fdt_ptr, buf + tag_pos, buf_size);
+
   // Check FDT header
-  const char *fdt_ptr = buf + tag_pos;
-  if (fdt_check_header(fdt_ptr) != 0) {
-    fprintf(stderr, "Bad DTB header.\n");
-    goto end;
+  if (fdt_check_full(fdt_ptr, buf_size) != 0) {
+    fprintf(stderr, "Bad DTB.\n");
+    goto error;
   }
 
   // Check FDT size and actual size
   size_t fdt_size = fdt_totalsize(fdt_ptr);
-  size_t fdt_actual_size = buf_size - tag_pos;
-  int fdt_size_diff = (int)fdt_actual_size - (int)fdt_size;
-  if (fdt_size_diff) {
-    fprintf(stderr, "Wrong size: actual size = %d FDT size = %d(%d)\n",
-      fdt_actual_size, fdt_size, fdt_size_diff);
+  if (buf_size < fdt_size) {
+    fprintf(stderr,
+            "Wrong size: fdt truncated: buffer size = %zu < FDT size = %zu\n",
+            buf_size, fdt_size);
+    goto error;
   }
 
   // Print the DT basic information
   int root_node_off = fdt_path_offset(fdt_ptr, "/");
   if (root_node_off < 0) {
     fprintf(stderr, "Can not get the root node.\n");
-    goto end;
+    goto error;
   }
   printf("Output %s\n", filename);
   const char *model =
@@ -84,16 +93,20 @@ int find_and_write_dtb(const char *filename,
   // Output DTB file
   if (write_fdt_to_file(filename, fdt_ptr) != 0) {
     fprintf(stderr, "Write file error: %s\n", filename);
-    goto end;
+    goto error;
   }
 
-end:
+  free(fdt_ptr);
+
   return tag_pos;
+
+error:
+  if (fdt_ptr) free(fdt_ptr);
+  return -1;
 }
 
-int extract_dtbs(const char *in_filename,
-                 const char *out_dtb_filename,
-                 const char *out_image_filename) {
+static int extract_dtbs(const char *in_filename, const char *out_dtb_filename,
+                        const char *out_image_filename) {
   int ret = 1;
   char *buf = NULL;