1 use std::{fs::File, io::Write, mem::ManuallyDrop, 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, skip_if_exists: bool) -> Result<()> { 9 let disk_image_path = ctx.disk_image_path(); 10 if disk_image_path.exists() { 11 if skip_if_exists { 12 return Ok(()); 13 } 14 return Err(anyhow!( 15 "Disk image already exists: {}", 16 disk_image_path.display() 17 )); 18 } 19 20 disk_path_safety_check(&disk_image_path)?; 21 22 // 获取镜像大小 23 let image_size = ctx.disk_image_size(); 24 create_raw_img(&disk_image_path, image_size).expect("Failed to create raw disk image"); 25 26 // 判断是否需要分区？ 27 28 let r = if ctx.rootfs().partition.image_should_be_partitioned() { 29 create_partitioned_image(ctx, &disk_image_path) 30 } else { 31 create_unpartitioned_image(ctx, &disk_image_path) 32 }; 33 34 if r.is_err() { 35 std::fs::remove_file(&disk_image_path).expect("Failed to remove disk image"); 36 } 37 r 38 } 39 40 pub(super) fn delete(ctx: &DADKExecContext, skip_if_not_exists: bool) -> Result<()> { 41 let disk_image_path = ctx.disk_image_path(); 42 if !disk_image_path.exists() { 43 if skip_if_not_exists { 44 return Ok(()); 45 } 46 return Err(anyhow!( 47 "Disk image does not exist: {}", 48 disk_image_path.display() 49 )); 50 } 51 disk_path_safety_check(&disk_image_path)?; 52 53 std::fs::remove_file(&disk_image_path) 54 .map_err(|e| anyhow!("Failed to remove disk image: {}", e))?; 55 Ok(()) 56 } 57 58 pub fn mount(ctx: &DADKExecContext) -> Result<()> { 59 let disk_image_path = ctx.disk_image_path(); 60 if !disk_image_path.exists() { 61 return Err(anyhow!( 62 "Disk image does not exist: {}", 63 disk_image_path.display() 64 )); 65 } 66 let disk_mount_path = ctx.disk_mount_path(); 67 68 // 尝试创建挂载点 69 std::fs::create_dir_all(&disk_mount_path) 70 .map_err(|e| anyhow!("Failed to create disk mount path: {}", e))?; 71 72 let partitioned = ctx.rootfs().partition.image_should_be_partitioned(); 73 log::trace!("Disk image is partitioned: {}", partitioned); 74 if partitioned { 75 mount_partitioned_image(ctx, &disk_image_path, &disk_mount_path)? 76 } else { 77 mount_unpartitioned_image(ctx, &disk_image_path, &disk_mount_path)? 78 } 79 log::info!("Disk image mounted at {}", disk_mount_path.display()); 80 Ok(()) 81 } 82 83 fn mount_partitioned_image( 84 ctx: &DADKExecContext, 85 disk_image_path: &PathBuf, 86 disk_mount_path: &PathBuf, 87 ) -> Result<()> { 88 let mut loop_device = ManuallyDrop::new( 89 LoopDeviceBuilder::new() 90 .img_path(disk_image_path.clone()) 91 .build() 92 .map_err(|e| anyhow!("Failed to create loop device: {}", e))?, 93 ); 94 95 loop_device 96 .attach() 97 .map_err(|e| anyhow!("Failed to attach loop device: {}", e))?; 98 99 let dev_path = loop_device.partition_path(1)?; 100 mount_unpartitioned_image(ctx, &dev_path, disk_mount_path)?; 101 102 Ok(()) 103 } 104 105 fn mount_unpartitioned_image( 106 _ctx: &DADKExecContext, 107 disk_image_path: &PathBuf, 108 disk_mount_path: &PathBuf, 109 ) -> Result<()> { 110 let cmd = Command::new("mount") 111 .arg(disk_image_path) 112 .arg(disk_mount_path) 113 .output() 114 .map_err(|e| anyhow!("Failed to mount disk image: {}", e))?; 115 if !cmd.status.success() { 116 return Err(anyhow!( 117 "Failed to mount disk image: {}", 118 String::from_utf8_lossy(&cmd.stderr) 119 )); 120 } 121 Ok(()) 122 } 123 124 pub fn umount(ctx: &DADKExecContext) -> Result<()> { 125 let disk_img_path = ctx.disk_image_path(); 126 let disk_mount_path = ctx.disk_mount_path(); 127 let mut loop_device = LoopDeviceBuilder::new().img_path(disk_img_path).build(); 128 129 let should_detach_loop_device: bool; 130 if let Ok(loop_device) = loop_device.as_mut() { 131 if let Err(e) = loop_device.attach_by_exists() { 132 log::trace!("umount: Failed to attach loop device: {}", e); 133 } 134 135 should_detach_loop_device = loop_device.attached(); 136 } else { 137 should_detach_loop_device = false; 138 } 139 140 if disk_mount_path.exists() { 141 let cmd = Command::new("umount") 142 .arg(disk_mount_path) 143 .output() 144 .map_err(|e| anyhow!("Failed to umount disk image: {}", e)); 145 match cmd { 146 Ok(cmd) => { 147 if !cmd.status.success() { 148 let e = anyhow!( 149 "Failed to umount disk image: {}", 150 String::from_utf8_lossy(&cmd.stderr) 151 ); 152 if should_detach_loop_device { 153 log::error!("{}", e); 154 } else { 155 return Err(e); 156 } 157 } 158 } 159 Err(e) => { 160 if should_detach_loop_device { 161 log::error!("{}", e); 162 } else { 163 return Err(e); 164 } 165 } 166 } 167 } 168 169 if let Ok(mut loop_device) = loop_device { 170 let loop_dev_path = loop_device.dev_path().cloned(); 171 if let Err(e) = loop_device.detach().map_err(|e| anyhow!("{}", e)) { 172 if ctx.rootfs().partition.image_should_be_partitioned() { 173 return Err(e); 174 } 175 } 176 log::info!("Loop device detached: {:?}", loop_dev_path); 177 } 178 179 Ok(()) 180 } 181 182 /// Ensures the provided disk image path is not a device node. 183 fn disk_path_safety_check(disk_image_path: &PathBuf) -> Result<()> { 184 const DONT_ALLOWED_PREFIX: [&str; 5] = 185 ["/dev/sd", "/dev/hd", "/dev/vd", "/dev/nvme", "/dev/mmcblk"]; 186 let path = disk_image_path.to_str().ok_or(anyhow!( 187 "disk path safety check failed: disk path is not valid utf-8" 188 ))?; 189 190 DONT_ALLOWED_PREFIX.iter().for_each(|prefix| { 191 if path.starts_with(prefix) { 192 panic!("disk path safety check failed: disk path is not allowed to be a device node(except loop dev)"); 193 } 194 }); 195 Ok(()) 196 } 197 198 fn create_partitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 199 let part_type = ctx.rootfs().partition.partition_type; 200 DiskPartitioner::create_partitioned_image(disk_image_path, part_type)?; 201 // 挂载loop设备 202 let mut loop_device = LoopDeviceBuilder::new() 203 .img_path(disk_image_path.clone()) 204 .build() 205 .map_err(|e| anyhow!("Failed to create loop device: {}", e))?; 206 loop_device 207 .attach() 208 .map_err(|e| anyhow!("Failed to attach loop device: {}", e))?; 209 210 let partition_path = loop_device.partition_path(1)?; 211 let fs_type = ctx.rootfs().metadata.fs_type; 212 DiskFormatter::format_disk(&partition_path, &fs_type)?; 213 loop_device.detach()?; 214 Ok(()) 215 } 216 217 fn create_unpartitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 218 // 直接对整块磁盘镜像进行格式化 219 let fs_type = ctx.rootfs().metadata.fs_type; 220 DiskFormatter::format_disk(disk_image_path, &fs_type) 221 } 222 223 /// 创建全0的raw镜像 224 fn create_raw_img(disk_image_path: &PathBuf, image_size: usize) -> Result<()> { 225 log::trace!("Creating raw disk image: {}", disk_image_path.display()); 226 // 创建父目录 227 if let Some(parent) = disk_image_path.parent() { 228 log::trace!("Creating parent directory: {}", parent.display()); 229 std::fs::create_dir_all(parent)?; 230 } 231 // 打开或创建文件 232 let mut file = File::create(disk_image_path)?; 233 234 // 将文件大小设置为指定大小 235 file.set_len(image_size.try_into().unwrap())?; 236 237 // 写入全0数据 238 let zero_buffer = vec![0u8; 4096]; // 4KB buffer for writing zeros 239 let mut remaining_size = image_size; 240 241 while remaining_size > 0 { 242 let write_size = std::cmp::min(remaining_size, zero_buffer.len()); 243 file.write_all(&zero_buffer[..write_size as usize])?; 244 remaining_size -= write_size; 245 } 246 247 Ok(()) 248 } 249 250 struct DiskPartitioner; 251 252 impl DiskPartitioner { 253 fn create_partitioned_image(disk_image_path: &PathBuf, part_type: PartitionType) -> Result<()> { 254 match part_type { 255 PartitionType::None => { 256 // This case should not be reached as we are in the partitioned image creation function 257 return Err(anyhow::anyhow!("Invalid partition type: None")); 258 } 259 PartitionType::Mbr => { 260 // Create MBR partitioned disk image 261 Self::create_mbr_partitioned_image(disk_image_path)?; 262 } 263 PartitionType::Gpt => { 264 // Create GPT partitioned disk image 265 Self::create_gpt_partitioned_image(disk_image_path)?; 266 } 267 } 268 Ok(()) 269 } 270 271 fn create_mbr_partitioned_image(disk_image_path: &PathBuf) -> Result<()> { 272 let disk_image_path_str = disk_image_path.to_str().expect("Invalid path"); 273 274 // 检查 fdisk 是否存在 275 let output = Command::new("fdisk") 276 .arg("--help") 277 .stdin(std::process::Stdio::piped()) 278 .stdout(std::process::Stdio::piped()) 279 .spawn()? 280 .wait_with_output()?; 281 282 if !output.status.success() { 283 return Err(anyhow::anyhow!("Command fdisk not found")); 284 } 285 286 // 向 fdisk 发送命令 287 let fdisk_commands = "o\nn\n\n\n\n\na\nw\n"; 288 let mut fdisk_child = Command::new("fdisk") 289 .arg(disk_image_path_str) 290 .stdin(std::process::Stdio::piped()) 291 .stdout(std::process::Stdio::piped()) 292 .spawn()?; 293 294 let fdisk_stdin = fdisk_child.stdin.as_mut().expect("Failed to open stdin"); 295 fdisk_stdin.write_all(fdisk_commands.as_bytes())?; 296 fdisk_stdin.flush()?; 297 fdisk_child 298 .wait() 299 .unwrap_or_else(|e| panic!("Failed to run fdisk: {}", e)); 300 Ok(()) 301 } 302 303 fn create_gpt_partitioned_image(_disk_image_path: &PathBuf) -> Result<()> { 304 // Implement the logic to create a GPT partitioned disk image 305 // This is a placeholder for the actual implementation 306 unimplemented!("Not implemented: create_gpt_partitioned_image"); 307 } 308 } 309 310 struct DiskFormatter; 311 312 impl DiskFormatter { 313 fn format_disk(disk_image_path: &PathBuf, fs_type: &FsType) -> Result<()> { 314 match fs_type { 315 FsType::Fat32 => Self::format_fat32(disk_image_path), 316 } 317 } 318 319 fn format_fat32(disk_image_path: &PathBuf) -> Result<()> { 320 // Use the `mkfs.fat` command to format the disk image as FAT32 321 let status = Command::new("mkfs.fat") 322 .arg("-F32") 323 .arg(disk_image_path.to_str().unwrap()) 324 .status()?; 325 326 if status.success() { 327 Ok(()) 328 } else { 329 Err(anyhow::anyhow!("Failed to format disk image as FAT32")) 330 } 331 } 332 } 333 334 #[cfg(test)] 335 mod tests { 336 use super::*; 337 use std::fs; 338 use std::io::Read; 339 use tempfile::NamedTempFile; 340 341 #[test] 342 fn test_create_raw_img_functional() -> Result<()> { 343 // 创建一个临时文件路径 344 let temp_file = NamedTempFile::new()?; 345 let disk_image_path = temp_file.path().to_path_buf(); 346 let disk_image_size = 1024 * 1024usize; 347 348 // 调用函数 349 create_raw_img(&disk_image_path, disk_image_size)?; 350 351 // 验证文件大小 352 let metadata = fs::metadata(&disk_image_path)?; 353 assert_eq!(metadata.len(), disk_image_size as u64); 354 355 // 验证文件内容是否全为0 356 let mut file = File::open(&disk_image_path)?; 357 let mut buffer = vec![0u8; 4096]; 358 let mut all_zeros = true; 359 360 while file.read(&mut buffer)? > 0 { 361 for byte in &buffer { 362 if *byte != 0 { 363 all_zeros = false; 364 break; 365 } 366 } 367 } 368 369 assert!(all_zeros, "File content is not all zeros"); 370 371 Ok(()) 372 } 373 374 #[test] 375 fn test_format_fat32() { 376 // Create a temporary file to use as the disk image 377 let temp_file = NamedTempFile::new().expect("Failed to create temp file"); 378 let disk_image_path = temp_file.path().to_path_buf(); 379 380 // 16MB 381 let image_size = 16 * 1024 * 1024usize; 382 create_raw_img(&disk_image_path, image_size).expect("Failed to create raw disk image"); 383 384 // Call the function to format the disk image 385 DiskFormatter::format_disk(&disk_image_path, &FsType::Fat32) 386 .expect("Failed to format disk image as FAT32"); 387 388 // Optionally, you can check if the disk image was actually formatted as FAT32 389 // by running a command to inspect the filesystem type 390 let output = Command::new("file") 391 .arg("-sL") 392 .arg(&disk_image_path) 393 .output() 394 .expect("Failed to execute 'file' command"); 395 396 let output_str = String::from_utf8_lossy(&output.stdout); 397 assert!( 398 output_str.contains("FAT (32 bit)"), 399 "Disk image is not formatted as FAT32" 400 ); 401 } 402 403 #[test] 404 fn test_create_mbr_partitioned_image() -> Result<()> { 405 // Create a temporary file to use as the disk image 406 let temp_file = NamedTempFile::new()?; 407 let disk_image_path = temp_file.path().to_path_buf(); 408 409 eprintln!("Disk image path: {:?}", disk_image_path); 410 // Create a raw disk image 411 let disk_image_size = 16 * 1024 * 1024usize; // 16MB 412 create_raw_img(&disk_image_path, disk_image_size)?; 413 414 // Call the function to create the MBR partitioned image 415 DiskPartitioner::create_mbr_partitioned_image(&disk_image_path)?; 416 417 // Verify the disk image has been correctly partitioned 418 let output = Command::new("fdisk") 419 .env("LANG", "C") // Set LANG to C to force English output 420 .env("LC_ALL", "C") // Set LC_ALL to C to force English output 421 .arg("-l") 422 .arg(&disk_image_path) 423 .output() 424 .expect("Failed to execute 'fdisk -l' command"); 425 426 let output_str = String::from_utf8_lossy(&output.stdout); 427 assert!( 428 output_str.contains("Disklabel type: dos"), 429 "Disk image does not have an MBR partition table" 430 ); 431 assert!( 432 output_str.contains("Start"), 433 "Disk image does not have a partition" 434 ); 435 436 Ok(()) 437 } 438 } 439