nyx_lite/
image_builder.rs

1use std::fs::{self, File};
2use std::io;
3use std::path::{Path, PathBuf};
4use std::process::{Command, Stdio};
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use thiserror::Error;
8
9#[derive(Debug, Error)]
10pub enum ImageBuilderError {
11    #[error("io error: {0}")]
12    Io(#[from] io::Error),
13    #[error("command `{0}` failed with status {1}")]
14    CommandFailed(String, std::process::ExitStatus),
15    #[error("command `{0}` returned empty output")]
16    CommandEmptyOutput(String),
17    #[error("root directory is empty or missing: {0}")]
18    MissingRootDir(PathBuf),
19}
20
21#[derive(Debug, Clone)]
22pub struct RootfsBuilder {
23    work_dir: PathBuf,
24}
25
26impl RootfsBuilder {
27    pub fn new(work_dir: impl Into<PathBuf>) -> Self {
28        Self {
29            work_dir: work_dir.into(),
30        }
31    }
32
33    pub fn build_from_dockerfile(
34        &self,
35        dockerfile: &Path,
36        context_dir: &Path,
37        output_img: &Path,
38        size_mib: Option<u64>,
39    ) -> Result<(), ImageBuilderError> {
40        let tag = format!("nyx-lite-rootfs-{}", unique_id());
41        run_status(
42            Command::new("docker")
43                .arg("build")
44                .arg("-f")
45                .arg(dockerfile)
46                .arg("-t")
47                .arg(&tag)
48                .arg(context_dir)
49                .stdout(Stdio::inherit())
50                .stderr(Stdio::inherit()),
51        )?;
52        let _image_guard = DockerImageGuard::new(tag.clone());
53
54        let container_id = run_output(
55            Command::new("docker")
56                .arg("create")
57                .arg(&tag)
58                .stdout(Stdio::piped())
59                .stderr(Stdio::piped()),
60        )?;
61        let _container_guard = DockerContainerGuard::new(container_id.clone());
62
63        let temp_root = TempDir::new_in(&self.work_dir, "nyx-lite-rootfs")?;
64        let tar_path = temp_root.path().join("rootfs.tar");
65        run_status(
66            Command::new("docker")
67                .arg("export")
68                .arg("-o")
69                .arg(&tar_path)
70                .arg(&container_id)
71                .stdout(Stdio::inherit())
72                .stderr(Stdio::inherit()),
73        )?;
74
75        let root_dir = temp_root.path().join("root");
76        fs::create_dir_all(&root_dir)?;
77        run_status(
78            Command::new("tar")
79                .arg("-xf")
80                .arg(&tar_path)
81                .arg("-C")
82                .arg(&root_dir)
83                .stdout(Stdio::inherit())
84                .stderr(Stdio::inherit()),
85        )?;
86
87        self.build_from_rootdir(&root_dir, output_img, size_mib)
88    }
89
90    pub fn build_from_rootdir(
91        &self,
92        root_dir: &Path,
93        output_img: &Path,
94        size_mib: Option<u64>,
95    ) -> Result<(), ImageBuilderError> {
96        if !root_dir.is_dir() {
97            return Err(ImageBuilderError::MissingRootDir(root_dir.to_path_buf()));
98        }
99
100        let estimated_mib = size_mib.unwrap_or_else(|| {
101            let bytes = dir_size(root_dir).unwrap_or(0);
102            estimate_mib(bytes)
103        });
104
105        if let Some(parent) = output_img.parent() {
106            fs::create_dir_all(parent)?;
107        }
108
109        let file = File::create(output_img)?;
110        file.set_len(estimated_mib.saturating_mul(1024 * 1024))?;
111        drop(file);
112
113        run_status(
114            Command::new("mke2fs")
115                .arg("-t")
116                .arg("ext4")
117                .arg("-F")
118                .arg("-L")
119                .arg("rootfs")
120                .arg("-d")
121                .arg(root_dir)
122                .arg(output_img)
123                .stdout(Stdio::inherit())
124                .stderr(Stdio::inherit()),
125        )
126    }
127}
128
129fn estimate_mib(bytes: u64) -> u64 {
130    let extra = bytes / 10;
131    let extra = extra.max(64 * 1024 * 1024);
132    let total = bytes.saturating_add(extra);
133    let mib = (total + (1024 * 1024 - 1)) / (1024 * 1024);
134    mib.max(256)
135}
136
137fn dir_size(root: &Path) -> io::Result<u64> {
138    let mut total: u64 = 0;
139    for entry in fs::read_dir(root)? {
140        let entry = entry?;
141        let path = entry.path();
142        let metadata = entry.metadata()?;
143        if metadata.is_dir() {
144            total = total.saturating_add(dir_size(&path)?);
145        } else {
146            total = total.saturating_add(metadata.len());
147        }
148    }
149    Ok(total)
150}
151
152fn run_status(command: &mut Command) -> Result<(), ImageBuilderError> {
153    let name = format!("{:?}", command);
154    let status = command.status()?;
155    if status.success() {
156        Ok(())
157    } else {
158        Err(ImageBuilderError::CommandFailed(name, status))
159    }
160}
161
162fn run_output(command: &mut Command) -> Result<String, ImageBuilderError> {
163    let name = format!("{:?}", command);
164    let output = command.output()?;
165    if !output.status.success() {
166        return Err(ImageBuilderError::CommandFailed(name, output.status));
167    }
168    let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
169    if stdout.is_empty() {
170        Err(ImageBuilderError::CommandEmptyOutput(name))
171    } else {
172        Ok(stdout)
173    }
174}
175
176fn unique_id() -> String {
177    let pid = std::process::id();
178    let now = SystemTime::now()
179        .duration_since(UNIX_EPOCH)
180        .unwrap_or_default()
181        .as_millis();
182    format!("{pid}-{now}")
183}
184
185struct TempDir {
186    path: PathBuf,
187}
188
189impl TempDir {
190    fn new_in(base: &Path, prefix: &str) -> Result<Self, ImageBuilderError> {
191        let mut path = base.to_path_buf();
192        path.push(format!("{prefix}-{}", unique_id()));
193        fs::create_dir_all(&path)?;
194        Ok(Self { path })
195    }
196
197    fn path(&self) -> &Path {
198        &self.path
199    }
200}
201
202impl Drop for TempDir {
203    fn drop(&mut self) {
204        let _ = fs::remove_dir_all(&self.path);
205    }
206}
207
208struct DockerContainerGuard {
209    id: String,
210}
211
212impl DockerContainerGuard {
213    fn new(id: String) -> Self {
214        Self { id }
215    }
216}
217
218impl Drop for DockerContainerGuard {
219    fn drop(&mut self) {
220        let _ = Command::new("docker")
221            .arg("rm")
222            .arg("-f")
223            .arg(&self.id)
224            .status();
225    }
226}
227
228struct DockerImageGuard {
229    tag: String,
230}
231
232impl DockerImageGuard {
233    fn new(tag: String) -> Self {
234        Self { tag }
235    }
236}
237
238impl Drop for DockerImageGuard {
239    fn drop(&mut self) {
240        let _ = Command::new("docker")
241            .arg("rmi")
242            .arg("-f")
243            .arg(&self.tag)
244            .status();
245    }
246}