2 files changed, 37 insertions, 6 deletions

diff --git a/Cargo.toml b/Cargo.toml
index d482ff0..c46cbcc 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,3 +11,5 @@ clap = { version = "4.0.29", features = ["derive"] }
 fatfs = "0.3.5"
 fscommon = "0.1.1"
 nix = { version = "0.26.1", features = ["ioctl"] }
+squashfs-ng = "0.1.2"
+temp-file = "0.1.7"
diff --git a/src/main.rs b/src/main.rs
index 83e9c97..2388166 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,8 @@ use anyhow::bail;
 use clap::Parser;
 use fatfs::{FatType, FormatVolumeOptions};
 use fscommon::StreamSlice;
+use squashfs_ng::write::{Source as SqsSource, SourceData as SqsSourceData, Writer as SqsWriter};
+use std::collections::HashMap;
 use std::fs::{File, OpenOptions};
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::os::unix::fs::PermissionsExt;
@@ -75,7 +77,7 @@ fn write_mbr_partition_table(file: &mut File, dev_size: u64) -> anyhow::Result<(
     file.write_all(SQUASHFS)?;
     file.write_all(INVALID_CHS)?;
     file.write_all(&(2048 + 256 * MiB / 512).to_le_bytes())?;
-    file.write_all(&(dev_size as u32 / 512 - 8192 - 256 * MiB / 512).to_le_bytes())?;
+    file.write_all(&(dev_size as u32 / 512 - 2048 - 256 * MiB / 512).to_le_bytes())?;
 
     // Partition 3 (unused)
     file.write_all(NOPART)?;
@@ -90,17 +92,19 @@ fn write_mbr_partition_table(file: &mut File, dev_size: u64) -> anyhow::Result<(
 }
 
 fn partition(file: &mut File, dev_size: u64) -> anyhow::Result<()> {
+    const ROOT_START: u64 = (2048 * 512 + 256 * MiB) as u64;
+    let root_end = ROOT_START + (dev_size as u32 - 2048 * 512 - 256 * MiB) as u64;
+
     write_mbr_partition_table(file, dev_size)?;
 
-    let mut boot_partition = StreamSlice::new(
-        file.try_clone()?,
-        2048 * 512,
-        (2048 * 512 + 256 * MiB - 1).into(),
-    )?;
+    let mut boot_partition = StreamSlice::new(file.try_clone()?, 2048 * 512, ROOT_START - 1)?;
+    let mut root_partition = StreamSlice::new(file.try_clone()?, ROOT_START, root_end)?;
 
     write_boot(&mut boot_partition)?;
     write_mbr(file)?;
 
+    write_root(&mut root_partition)?;
+
     Ok(())
 }
 
@@ -194,6 +198,31 @@ fn write_mbr(file: &mut File) -> anyhow::Result<()> {
     Ok(())
 }
 
+fn write_root(partition: &mut StreamSlice<File>) -> anyhow::Result<()> {
+    let mut partition_buf = Vec::new();
+    partition.read_to_end(&mut partition_buf)?;
+
+    let tmp_file = temp_file::with_contents(&partition_buf);
+
+    let mut writer = SqsWriter::open(tmp_file.path())?;
+
+    let init_file = File::open("init")?;
+    writer.add(SqsSource {
+        data: SqsSourceData::File(Box::new(init_file)),
+        uid: 0,
+        gid: 0,
+        mode: 0o755,
+        modified: 0,
+        xattrs: HashMap::new(),
+        flags: 0,
+    })?;
+
+    writer.finish()?;
+
+    println!("Root filesystem created successfully");
+    Ok(())
+}
+
 fn overwrite_device(file: &mut File, overwrite: String) -> anyhow::Result<()> {
     partition_device(file, overwrite)?;
     Ok(())