1 use std::{fs::File, io::Write, path::PathBuf, process::Command}; 2 3 use crate::context::DADKExecContext; 4 use anyhow::{anyhow, Result}; 5 use dadk_config::rootfs::{fstype::FsType, partition::PartitionType}; 6 7 use super::loopdev::LoopDeviceBuilder; 8 pub(super) fn create(ctx: &DADKExecContext) -> Result<()> { 9 let disk_image_path = ctx.disk_image_path(); 10 if disk_image_path.exists() { 11 return Err(anyhow!( 12 "Disk image already exists: {}", 13 disk_image_path.display() 14 )); 15 } 16 17 disk_path_safety_check(&disk_image_path)?; 18 19 // 获取镜像大小 20 let image_size = ctx.disk_image_size(); 21 create_raw_img(&disk_image_path, image_size).expect("Failed to create raw disk image"); 22 23 // 判断是否需要分区？ 24 25 let r = if ctx.rootfs().partition.should_create_partitioned_image() { 26 create_partitioned_image(ctx, &disk_image_path) 27 } else { 28 create_unpartitioned_image(ctx, &disk_image_path) 29 }; 30 31 if r.is_err() { 32 std::fs::remove_file(&disk_image_path).expect("Failed to remove disk image"); 33 } 34 r 35 } 36 /// Ensures the provided disk image path is not a device node. 37 fn disk_path_safety_check(disk_image_path: &PathBuf) -> Result<()> { 38 const DONT_ALLOWED_PREFIX: [&str; 5] = 39 ["/dev/sd", "/dev/hd", "/dev/vd", "/dev/nvme", "/dev/mmcblk"]; 40 let path = disk_image_path.to_str().ok_or(anyhow!( 41 "disk path safety check failed: disk path is not valid utf-8" 42 ))?; 43 44 DONT_ALLOWED_PREFIX.iter().for_each(|prefix| { 45 if path.starts_with(prefix) { 46 panic!("disk path safety check failed: disk path is not allowed to be a device node(except loop dev)"); 47 } 48 }); 49 Ok(()) 50 } 51 52 fn create_partitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 53 let part_type = ctx.rootfs().partition.partition_type; 54 DiskPartitioner::create_partitioned_image(disk_image_path, part_type)?; 55 // 挂载loop设备 56 let mut loop_device = LoopDeviceBuilder::new() 57 .img_path(disk_image_path.clone()) 58 .build() 59 .map_err(|e| anyhow!("Failed to create loop device: {}", e))?; 60 61 let partition_path = loop_device.partition_path(1)?; 62 let fs_type = ctx.rootfs().metadata.fs_type; 63 DiskFormatter::format_disk(&partition_path, &fs_type)?; 64 loop_device.detach()?; 65 Ok(()) 66 } 67 68 fn create_unpartitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 69 // 直接对整块磁盘镜像进行格式化 70 let fs_type = ctx.rootfs().metadata.fs_type; 71 DiskFormatter::format_disk(disk_image_path, &fs_type) 72 } 73 74 /// 创建全0的raw镜像 75 fn create_raw_img(disk_image_path: &PathBuf, image_size: usize) -> Result<()> { 76 log::trace!("Creating raw disk image: {}", disk_image_path.display()); 77 // 创建父目录 78 if let Some(parent) = disk_image_path.parent() { 79 log::trace!("Creating parent directory: {}", parent.display()); 80 std::fs::create_dir_all(parent)?; 81 } 82 // 打开或创建文件 83 let mut file = File::create(disk_image_path)?; 84 85 // 将文件大小设置为指定大小 86 file.set_len(image_size.try_into().unwrap())?; 87 88 // 写入全0数据 89 let zero_buffer = vec![0u8; 4096]; // 4KB buffer for writing zeros 90 let mut remaining_size = image_size; 91 92 while remaining_size > 0 { 93 let write_size = std::cmp::min(remaining_size, zero_buffer.len()); 94 file.write_all(&zero_buffer[..write_size as usize])?; 95 remaining_size -= write_size; 96 } 97 98 Ok(()) 99 } 100 101 struct DiskPartitioner; 102 103 impl DiskPartitioner { 104 fn create_partitioned_image(disk_image_path: &PathBuf, part_type: PartitionType) -> Result<()> { 105 match part_type { 106 PartitionType::None => { 107 // This case should not be reached as we are in the partitioned image creation function 108 return Err(anyhow::anyhow!("Invalid partition type: None")); 109 } 110 PartitionType::Mbr => { 111 // Create MBR partitioned disk image 112 Self::create_mbr_partitioned_image(disk_image_path)?; 113 } 114 PartitionType::Gpt => { 115 // Create GPT partitioned disk image 116 Self::create_gpt_partitioned_image(disk_image_path)?; 117 } 118 } 119 Ok(()) 120 } 121 122 fn create_mbr_partitioned_image(disk_image_path: &PathBuf) -> Result<()> { 123 let disk_image_path_str = disk_image_path.to_str().expect("Invalid path"); 124 125 // 检查 fdisk 是否存在 126 let output = Command::new("fdisk") 127 .arg("--help") 128 .stdin(std::process::Stdio::piped()) 129 .stdout(std::process::Stdio::piped()) 130 .spawn()? 131 .wait_with_output()?; 132 133 if !output.status.success() { 134 return Err(anyhow::anyhow!("Command fdisk not found")); 135 } 136 137 // 向 fdisk 发送命令 138 let fdisk_commands = "o\nn\n\n\n\n\na\nw\n"; 139 let mut fdisk_child = Command::new("fdisk") 140 .arg(disk_image_path_str) 141 .stdin(std::process::Stdio::piped()) 142 .stdout(std::process::Stdio::piped()) 143 .spawn()?; 144 145 let fdisk_stdin = fdisk_child.stdin.as_mut().expect("Failed to open stdin"); 146 fdisk_stdin.write_all(fdisk_commands.as_bytes())?; 147 fdisk_stdin.flush()?; 148 fdisk_child 149 .wait() 150 .unwrap_or_else(|e| panic!("Failed to run fdisk: {}", e)); 151 Ok(()) 152 } 153 154 fn create_gpt_partitioned_image(_disk_image_path: &PathBuf) -> Result<()> { 155 // Implement the logic to create a GPT partitioned disk image 156 // This is a placeholder for the actual implementation 157 unimplemented!("Not implemented: create_gpt_partitioned_image"); 158 } 159 } 160 161 struct DiskFormatter; 162 163 impl DiskFormatter { 164 fn format_disk(disk_image_path: &PathBuf, fs_type: &FsType) -> Result<()> { 165 match fs_type { 166 FsType::Fat32 => Self::format_fat32(disk_image_path), 167 } 168 } 169 170 fn format_fat32(disk_image_path: &PathBuf) -> Result<()> { 171 // Use the `mkfs.fat` command to format the disk image as FAT32 172 let status = Command::new("mkfs.fat") 173 .arg("-F32") 174 .arg(disk_image_path.to_str().unwrap()) 175 .status()?; 176 177 if status.success() { 178 Ok(()) 179 } else { 180 Err(anyhow::anyhow!("Failed to format disk image as FAT32")) 181 } 182 } 183 } 184 185 #[cfg(test)] 186 mod tests { 187 use super::*; 188 use std::fs; 189 use std::io::Read; 190 use tempfile::NamedTempFile; 191 192 #[test] 193 fn test_create_raw_img_functional() -> Result<()> { 194 // 创建一个临时文件路径 195 let temp_file = NamedTempFile::new()?; 196 let disk_image_path = temp_file.path().to_path_buf(); 197 let disk_image_size = 1024 * 1024usize; 198 199 // 调用函数 200 create_raw_img(&disk_image_path, disk_image_size)?; 201 202 // 验证文件大小 203 let metadata = fs::metadata(&disk_image_path)?; 204 assert_eq!(metadata.len(), disk_image_size as u64); 205 206 // 验证文件内容是否全为0 207 let mut file = File::open(&disk_image_path)?; 208 let mut buffer = vec![0u8; 4096]; 209 let mut all_zeros = true; 210 211 while file.read(&mut buffer)? > 0 { 212 for byte in &buffer { 213 if *byte != 0 { 214 all_zeros = false; 215 break; 216 } 217 } 218 } 219 220 assert!(all_zeros, "File content is not all zeros"); 221 222 Ok(()) 223 } 224 225 #[test] 226 fn test_format_fat32() { 227 // Create a temporary file to use as the disk image 228 let temp_file = NamedTempFile::new().expect("Failed to create temp file"); 229 let disk_image_path = temp_file.path().to_path_buf(); 230 231 // 16MB 232 let image_size = 16 * 1024 * 1024usize; 233 create_raw_img(&disk_image_path, image_size).expect("Failed to create raw disk image"); 234 235 // Call the function to format the disk image 236 DiskFormatter::format_disk(&disk_image_path, &FsType::Fat32) 237 .expect("Failed to format disk image as FAT32"); 238 239 // Optionally, you can check if the disk image was actually formatted as FAT32 240 // by running a command to inspect the filesystem type 241 let output = Command::new("file") 242 .arg("-sL") 243 .arg(&disk_image_path) 244 .output() 245 .expect("Failed to execute 'file' command"); 246 247 let output_str = String::from_utf8_lossy(&output.stdout); 248 assert!( 249 output_str.contains("FAT (32 bit)"), 250 "Disk image is not formatted as FAT32" 251 ); 252 } 253 254 #[test] 255 fn test_create_mbr_partitioned_image() -> Result<()> { 256 // Create a temporary file to use as the disk image 257 let temp_file = NamedTempFile::new()?; 258 let disk_image_path = temp_file.path().to_path_buf(); 259 260 eprintln!("Disk image path: {:?}", disk_image_path); 261 // Create a raw disk image 262 let disk_image_size = 16 * 1024 * 1024usize; // 16MB 263 create_raw_img(&disk_image_path, disk_image_size)?; 264 265 // Call the function to create the MBR partitioned image 266 DiskPartitioner::create_mbr_partitioned_image(&disk_image_path)?; 267 268 // Verify the disk image has been correctly partitioned 269 let output = Command::new("fdisk") 270 .env("LANG", "C") // Set LANG to C to force English output 271 .env("LC_ALL", "C") // Set LC_ALL to C to force English output 272 .arg("-l") 273 .arg(&disk_image_path) 274 .output() 275 .expect("Failed to execute 'fdisk -l' command"); 276 277 let output_str = String::from_utf8_lossy(&output.stdout); 278 assert!( 279 output_str.contains("Disklabel type: dos"), 280 "Disk image does not have an MBR partition table" 281 ); 282 assert!( 283 output_str.contains("Start"), 284 "Disk image does not have a partition" 285 ); 286 287 Ok(()) 288 } 289 } 290