Merge pull request #345 from Dstack-TEE/fix-overflow

dstack-mr: Fix potential panic due to int overflow

Author: kvinwang

dstack-mr/src/kernel.rs

@@ -99,7 +99,7 @@ fn authenticode_sha384_hash(data: &[u8]) -> Result<Vec<u8>> {
         let trailing_data_len = file_size - sum_of_bytes_hashed;
 
         if trailing_data_len > cert_table_size {
-            let hashed_trailing_len = trailing_data_len - cert_table_size;
+            let hashed_trailing_len = trailing_data_len.saturating_sub(cert_table_size);
             let trailing_start = sum_of_bytes_hashed;
 
             if trailing_start + hashed_trailing_len <= data.len() {
@@ -142,14 +142,14 @@ fn patch_kernel(
     }
     if protocol >= 0x201 {
         kd[0x211] |= 0x80; // loadflags |= CAN_USE_HEAP
-        let heap_end_ptr = cmdline_addr - real_addr - 0x200;
+        let heap_end_ptr = cmdline_addr.saturating_sub(real_addr).saturating_sub(0x200);
         kd[0x224..0x228].copy_from_slice(&heap_end_ptr.to_le_bytes());
     }
     if protocol >= 0x202 {
         kd[0x228..0x22C].copy_from_slice(&cmdline_addr.to_le_bytes());
     } else {
         kd[0x20..0x22].copy_from_slice(&0xa33f_u16.to_le_bytes());
-        let offset = (cmdline_addr - real_addr) as u16;
+        let offset = cmdline_addr.saturating_sub(real_addr) as u16;
         kd[0x22..0x24].copy_from_slice(&offset.to_le_bytes());
     }
 
@@ -186,14 +186,23 @@ fn patch_kernel(
             mem_size as u32
         };
 
-        if initrd_max >= below_4g_mem_size - acpi_data_size {
-            initrd_max = below_4g_mem_size - acpi_data_size - 1;
+        if let Some(available_mem) = below_4g_mem_size.checked_sub(acpi_data_size) {
+            if initrd_max >= available_mem {
+                initrd_max = available_mem.saturating_sub(1);
+            }
+        } else {
+            // If acpi_data_size >= below_4g_mem_size, we have no memory available
+            bail!(
+                "ACPI data size ({}) exceeds available memory ({})",
+                acpi_data_size,
+                below_4g_mem_size
+            );
         }
         if initrd_size >= initrd_max {
             bail!("initrd is too large");
         }
 
-        let initrd_addr = (initrd_max - initrd_size) & !4095;
+        let initrd_addr = initrd_max.saturating_sub(initrd_size) & !4095;
         kd[0x218..0x21C].copy_from_slice(&initrd_addr.to_le_bytes());
         kd[0x21C..0x220].copy_from_slice(&initrd_size.to_le_bytes());
     }

dstack-mr/src/tdvf.rs

@@ -82,20 +82,30 @@ impl<'a> Tdvf<'a> {
         const TABLE_FOOTER_GUID: &str = "96b582de-1fb2-45f7-baea-a366c55a082d";
         const BYTES_AFTER_TABLE_FOOTER: usize = 32;
 
+        if fw.len() < BYTES_AFTER_TABLE_FOOTER {
+            bail!("TDVF firmware too small");
+        }
         let offset = fw.len() - BYTES_AFTER_TABLE_FOOTER;
         let encoded_footer_guid = encode_guid(TABLE_FOOTER_GUID)?;
+        if offset < 16 {
+            bail!("TDVF firmware offset too small for GUID");
+        }
         let guid = &fw[offset - 16..offset];
 
         if guid != encoded_footer_guid {
             bail!("Failed to parse TDVF metadata: Invalid footer GUID");
         }
 
+        if offset < 18 {
+            bail!("TDVF firmware offset too small for tables length");
+        }
         let tables_len =
             u16::from_le_bytes(fw[offset - 18..offset - 16].try_into().unwrap()) as usize;
-        if tables_len == 0 || tables_len > offset - 18 {
+        if tables_len == 0 || tables_len > offset.saturating_sub(18) {
             bail!("Failed to parse TDVF metadata: Invalid tables length");
         }
-        let tables = &fw[offset - 18 - tables_len..offset - 18];
+        let table_start = offset.saturating_sub(18).saturating_sub(tables_len);
+        let tables = &fw[table_start..offset - 18];
         let mut offset = tables.len();
 
         let mut data: Option<&[u8]> = None;
@@ -106,21 +116,28 @@ impl<'a> Tdvf<'a> {
             }
             let guid = &tables[offset - 16..offset];
             let entry_len = read_le::<u16>(tables, offset - 18, "entry length")? as usize;
-            if entry_len > offset - 18 {
+            if entry_len > offset.saturating_sub(18) {
                 bail!("Failed to parse TDVF metadata: Invalid entry length");
             }
             if guid == encoded_guid {
-                data = Some(&tables[offset - 18 - entry_len..offset - 18]);
+                let entry_start = offset.saturating_sub(18).saturating_sub(entry_len);
+                data = Some(&tables[entry_start..offset - 18]);
                 break;
             }
-            offset -= entry_len;
+            offset = offset.saturating_sub(entry_len);
         }
 
         let data = data.context("Failed to parse TDVF metadata: Missing TDVF metadata")?;
 
-        let tdvf_meta_offset =
+        if data.len() < 4 {
+            bail!("TDVF metadata data too small");
+        }
+        let tdvf_meta_offset_raw =
             u32::from_le_bytes(data[data.len() - 4..].try_into().unwrap()) as usize;
-        let tdvf_meta_offset = fw.len() - tdvf_meta_offset;
+        if tdvf_meta_offset_raw > fw.len() {
+            bail!("TDVF metadata offset exceeds firmware size");
+        }
+        let tdvf_meta_offset = fw.len() - tdvf_meta_offset_raw;
         let tdvf_meta_desc = &fw[tdvf_meta_offset..tdvf_meta_offset + 16];
 
         if &tdvf_meta_desc[..4] != b"TDVF" {
@@ -311,16 +328,27 @@ impl<'a> Tdvf<'a> {
         let (_, last_start, last_end) = memory_acceptor.ranges.pop().expect("No ranges");
 
         for (accepted, start, end) in memory_acceptor.ranges {
+            if end < start {
+                bail!("Invalid memory range: end < start");
+            }
+            let size = end - start;
             if accepted {
-                add_memory_resource_hob(0x00, start, end - start);
+                add_memory_resource_hob(0x00, start, size);
             } else {
-                add_memory_resource_hob(0x07, start, end - start);
+                add_memory_resource_hob(0x07, start, size);
             }
         }
 
+        if last_end < last_start {
+            bail!("Invalid last memory range: end < start");
+        }
         if memory_size >= 0xB0000000 {
-            add_memory_resource_hob(0x07, last_start, 0x80000000u64 - last_start);
-            add_memory_resource_hob(0x07, 0x100000000, last_end - 0x80000000u64);
+            if last_start < 0x80000000u64 {
+                add_memory_resource_hob(0x07, last_start, 0x80000000u64 - last_start);
+            }
+            if last_end > 0x80000000u64 {
+                add_memory_resource_hob(0x07, 0x100000000, last_end - 0x80000000u64);
+            }
         } else {
             add_memory_resource_hob(0x07, last_start, last_end - last_start);
         }