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 loop_device.detach().ok(); 172 173 log::info!("Loop device detached: {:?}", loop_dev_path); 174 } 175 176 Ok(()) 177 } 178 179 /// Ensures the provided disk image path is not a device node. 180 fn disk_path_safety_check(disk_image_path: &PathBuf) -> Result<()> { 181 const DONT_ALLOWED_PREFIX: [&str; 5] = 182 ["/dev/sd", "/dev/hd", "/dev/vd", "/dev/nvme", "/dev/mmcblk"]; 183 let path = disk_image_path.to_str().ok_or(anyhow!( 184 "disk path safety check failed: disk path is not valid utf-8" 185 ))?; 186 187 DONT_ALLOWED_PREFIX.iter().for_each(|prefix| { 188 if path.starts_with(prefix) { 189 panic!("disk path safety check failed: disk path is not allowed to be a device node(except loop dev)"); 190 } 191 }); 192 Ok(()) 193 } 194 195 fn create_partitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 196 let part_type = ctx.rootfs().partition.partition_type; 197 DiskPartitioner::create_partitioned_image(disk_image_path, part_type)?; 198 // 挂载loop设备 199 let mut loop_device = LoopDeviceBuilder::new() 200 .img_path(disk_image_path.clone()) 201 .build() 202 .map_err(|e| anyhow!("Failed to create loop device: {}", e))?; 203 loop_device 204 .attach() 205 .map_err(|e| anyhow!("Failed to attach loop device: {}", e))?; 206 207 let partition_path = loop_device.partition_path(1)?; 208 let fs_type = ctx.rootfs().metadata.fs_type; 209 DiskFormatter::format_disk(&partition_path, &fs_type)?; 210 loop_device.detach()?; 211 Ok(()) 212 } 213 214 fn create_unpartitioned_image(ctx: &DADKExecContext, disk_image_path: &PathBuf) -> Result<()> { 215 // 直接对整块磁盘镜像进行格式化 216 let fs_type = ctx.rootfs().metadata.fs_type; 217 DiskFormatter::format_disk(disk_image_path, &fs_type) 218 } 219 220 /// 创建全0的raw镜像 221 fn create_raw_img(disk_image_path: &PathBuf, image_size: usize) -> Result<()> { 222 log::trace!("Creating raw disk image: {}", disk_image_path.display()); 223 // 创建父目录 224 if let Some(parent) = disk_image_path.parent() { 225 log::trace!("Creating parent directory: {}", parent.display()); 226 std::fs::create_dir_all(parent)?; 227 } 228 // 打开或创建文件 229 let mut file = File::create(disk_image_path)?; 230 231 // 将文件大小设置为指定大小 232 file.set_len(image_size.try_into().unwrap())?; 233 234 // 写入全0数据 235 let zero_buffer = vec![0u8; 4096]; // 4KB buffer for writing zeros 236 let mut remaining_size = image_size; 237 238 while remaining_size > 0 { 239 let write_size = std::cmp::min(remaining_size, zero_buffer.len()); 240 file.write_all(&zero_buffer[..write_size as usize])?; 241 remaining_size -= write_size; 242 } 243 244 Ok(()) 245 } 246 247 pub fn check_disk_image_exists(ctx: &DADKExecContext) -> Result<()> { 248 let disk_image_path = ctx.disk_image_path(); 249 if disk_image_path.exists() { 250 println!("1"); 251 } else { 252 println!("0"); 253 } 254 Ok(()) 255 } 256 257 pub fn show_mount_point(ctx: &DADKExecContext) -> Result<()> { 258 let disk_mount_path = ctx.disk_mount_path(); 259 println!("{}", disk_mount_path.display()); 260 Ok(()) 261 } 262 263 pub fn show_loop_device(ctx: &DADKExecContext) -> Result<()> { 264 let disk_image_path = ctx.disk_image_path(); 265 let mut loop_device = LoopDeviceBuilder::new().img_path(disk_image_path).build()?; 266 if let Err(e) = loop_device.attach_by_exists() { 267 log::error!("Failed to attach loop device: {}", e); 268 } else { 269 println!("{}", loop_device.dev_path().unwrap()); 270 } 271 Ok(()) 272 } 273 274 struct DiskPartitioner; 275 276 impl DiskPartitioner { 277 fn create_partitioned_image(disk_image_path: &PathBuf, part_type: PartitionType) -> Result<()> { 278 match part_type { 279 PartitionType::None => { 280 // This case should not be reached as we are in the partitioned image creation function 281 return Err(anyhow::anyhow!("Invalid partition type: None")); 282 } 283 PartitionType::Mbr => { 284 // Create MBR partitioned disk image 285 Self::create_mbr_partitioned_image(disk_image_path)?; 286 } 287 PartitionType::Gpt => { 288 // Create GPT partitioned disk image 289 Self::create_gpt_partitioned_image(disk_image_path)?; 290 } 291 } 292 Ok(()) 293 } 294 295 fn create_mbr_partitioned_image(disk_image_path: &PathBuf) -> Result<()> { 296 let disk_image_path_str = disk_image_path.to_str().expect("Invalid path"); 297 298 // 检查 fdisk 是否存在 299 let output = Command::new("fdisk") 300 .arg("--help") 301 .stdin(std::process::Stdio::piped()) 302 .stdout(std::process::Stdio::piped()) 303 .spawn()? 304 .wait_with_output()?; 305 306 if !output.status.success() { 307 return Err(anyhow::anyhow!("Command fdisk not found")); 308 } 309 310 // 向 fdisk 发送命令 311 let fdisk_commands = "o\nn\n\n\n\n\na\nw\n"; 312 let mut fdisk_child = Command::new("fdisk") 313 .arg(disk_image_path_str) 314 .stdin(std::process::Stdio::piped()) 315 .stdout(std::process::Stdio::piped()) 316 .spawn()?; 317 318 let fdisk_stdin = fdisk_child.stdin.as_mut().expect("Failed to open stdin"); 319 fdisk_stdin.write_all(fdisk_commands.as_bytes())?; 320 fdisk_stdin.flush()?; 321 fdisk_child 322 .wait() 323 .unwrap_or_else(|e| panic!("Failed to run fdisk: {}", e)); 324 Ok(()) 325 } 326 327 fn create_gpt_partitioned_image(_disk_image_path: &PathBuf) -> Result<()> { 328 // Implement the logic to create a GPT partitioned disk image 329 // This is a placeholder for the actual implementation 330 unimplemented!("Not implemented: create_gpt_partitioned_image"); 331 } 332 } 333 334 struct DiskFormatter; 335 336 impl DiskFormatter { 337 fn format_disk(disk_image_path: &PathBuf, fs_type: &FsType) -> Result<()> { 338 match fs_type { 339 FsType::Fat32 => Self::format_fat32(disk_image_path), 340 } 341 } 342 343 fn format_fat32(disk_image_path: &PathBuf) -> Result<()> { 344 // Use the `mkfs.fat` command to format the disk image as FAT32 345 let status = Command::new("mkfs.fat") 346 .arg("-F32") 347 .arg(disk_image_path.to_str().unwrap()) 348 .status()?; 349 350 if status.success() { 351 Ok(()) 352 } else { 353 Err(anyhow::anyhow!("Failed to format disk image as FAT32")) 354 } 355 } 356 } 357 358 #[cfg(test)] 359 mod tests { 360 use super::*; 361 use std::fs; 362 use std::io::Read; 363 use tempfile::NamedTempFile; 364 365 #[test] 366 fn test_create_raw_img_functional() -> Result<()> { 367 // 创建一个临时文件路径 368 let temp_file = NamedTempFile::new()?; 369 let disk_image_path = temp_file.path().to_path_buf(); 370 let disk_image_size = 1024 * 1024usize; 371 372 // 调用函数 373 create_raw_img(&disk_image_path, disk_image_size)?; 374 375 // 验证文件大小 376 let metadata = fs::metadata(&disk_image_path)?; 377 assert_eq!(metadata.len(), disk_image_size as u64); 378 379 // 验证文件内容是否全为0 380 let mut file = File::open(&disk_image_path)?; 381 let mut buffer = vec![0u8; 4096]; 382 let mut all_zeros = true; 383 384 while file.read(&mut buffer)? > 0 { 385 for byte in &buffer { 386 if *byte != 0 { 387 all_zeros = false; 388 break; 389 } 390 } 391 } 392 393 assert!(all_zeros, "File content is not all zeros"); 394 395 Ok(()) 396 } 397 398 #[test] 399 fn test_format_fat32() { 400 // Create a temporary file to use as the disk image 401 let temp_file = NamedTempFile::new().expect("Failed to create temp file"); 402 let disk_image_path = temp_file.path().to_path_buf(); 403 404 // 16MB 405 let image_size = 16 * 1024 * 1024usize; 406 create_raw_img(&disk_image_path, image_size).expect("Failed to create raw disk image"); 407 408 // Call the function to format the disk image 409 DiskFormatter::format_disk(&disk_image_path, &FsType::Fat32) 410 .expect("Failed to format disk image as FAT32"); 411 412 // Optionally, you can check if the disk image was actually formatted as FAT32 413 // by running a command to inspect the filesystem type 414 let output = Command::new("file") 415 .arg("-sL") 416 .arg(&disk_image_path) 417 .output() 418 .expect("Failed to execute 'file' command"); 419 420 let output_str = String::from_utf8_lossy(&output.stdout); 421 assert!( 422 output_str.contains("FAT (32 bit)"), 423 "Disk image is not formatted as FAT32" 424 ); 425 } 426 427 #[test] 428 fn test_create_mbr_partitioned_image() -> Result<()> { 429 // Create a temporary file to use as the disk image 430 let temp_file = NamedTempFile::new()?; 431 let disk_image_path = temp_file.path().to_path_buf(); 432 433 eprintln!("Disk image path: {:?}", disk_image_path); 434 // Create a raw disk image 435 let disk_image_size = 16 * 1024 * 1024usize; // 16MB 436 create_raw_img(&disk_image_path, disk_image_size)?; 437 438 // Call the function to create the MBR partitioned image 439 DiskPartitioner::create_mbr_partitioned_image(&disk_image_path)?; 440 441 // Verify the disk image has been correctly partitioned 442 let output = Command::new("fdisk") 443 .env("LANG", "C") // Set LANG to C to force English output 444 .env("LC_ALL", "C") // Set LC_ALL to C to force English output 445 .arg("-l") 446 .arg(&disk_image_path) 447 .output() 448 .expect("Failed to execute 'fdisk -l' command"); 449 450 let output_str = String::from_utf8_lossy(&output.stdout); 451 assert!( 452 output_str.contains("Disklabel type: dos"), 453 "Disk image does not have an MBR partition table" 454 ); 455 assert!( 456 output_str.contains("Start"), 457 "Disk image does not have a partition" 458 ); 459 460 Ok(()) 461 } 462 } 463