include_fs/
lib.rs

1use std::collections::HashMap;
2use std::env;
3use std::fs::File;
4use std::io::{self, Write};
5use std::path::{Path, PathBuf};
6use std::sync::LazyLock;
7use thiserror::Error;
8use walkdir::WalkDir;
9
10pub use include_fs_macros::include_fs;
11
12const MAGIC: &[u8; 4] = b"INFS";
13
14#[derive(Error, Debug)]
15pub enum ArchiveError {
16  #[error("Path too long: {path} ({len} bytes, max {max} bytes)")]
17  PathTooLong {
18    path: String,
19    len: usize,
20    max: usize,
21  },
22
23  #[error("Too many files: {count} (max {max})")]
24  TooManyFiles { count: usize, max: usize },
25
26  #[error("IO error: {0}")]
27  Io(#[from] std::io::Error),
28
29  #[error("Source directory must be a subdirectory of the manifest directory")]
30  InvalidSourceDirectory,
31
32  #[error("Failed to collect files: {0}")]
33  WalkDir(#[from] walkdir::Error),
34}
35
36#[derive(Error, Debug)]
37pub enum FsError {
38  #[error("File not found")]
39  NotFound,
40
41  #[error("Invalid archive")]
42  InvalidArchive,
43}
44
45#[derive(Debug)]
46struct FileEntry {
47  pub path: PathBuf,
48  pub size: u64,
49}
50
51impl FileEntry {
52  pub fn new(path: impl Into<PathBuf>, size: u64) -> Self {
53    Self {
54      path: path.into(),
55      size,
56    }
57  }
58}
59
60fn compute_header(files: &[FileEntry]) -> Result<Vec<u8>, ArchiveError> {
61  // Validate file count fits in u32
62  if files.len() > u32::MAX as usize {
63    return Err(ArchiveError::TooManyFiles {
64      count: files.len(),
65      max: u32::MAX as usize,
66    });
67  }
68
69  let mut header_size = 4 + 4; // magic + file count
70  for file in files {
71    let path_str = file.path.to_string_lossy();
72    let path_len = path_str.len();
73
74    if path_len > u16::MAX as usize {
75      return Err(ArchiveError::PathTooLong {
76        path: path_str.to_string(),
77        len: path_str.len(),
78        max: u16::MAX as usize,
79      });
80    }
81
82    // path_len + path + size + offset
83    header_size += 2 + path_len + 8 + 8;
84  }
85
86  let mut header = Vec::with_capacity(header_size);
87
88  header.extend_from_slice(MAGIC);
89  header.extend_from_slice(&(files.len() as u32).to_le_bytes());
90
91  let mut data_offset = header_size as u64;
92  for file in files {
93    let path_str = file.path.to_string_lossy();
94    let path_bytes = path_str.as_bytes();
95
96    header.extend_from_slice(&(path_bytes.len() as u16).to_le_bytes());
97    header.extend_from_slice(path_bytes);
98    header.extend_from_slice(&file.size.to_le_bytes());
99    header.extend_from_slice(&data_offset.to_le_bytes());
100
101    data_offset += file.size;
102  }
103
104  Ok(header)
105}
106
107fn write_archive(files: &[FileEntry], output_path: &Path) -> Result<(), ArchiveError> {
108  let mut file = File::create(output_path)?;
109
110  // Write header
111  let header = compute_header(files)?;
112  file.write_all(&header)?;
113
114  // Write file data
115  for file_entry in files {
116    let mut f = File::open(&file_entry.path)?;
117    io::copy(&mut f, &mut file)?;
118  }
119
120  Ok(())
121}
122
123/// Bundle a directory to be embedded in the binary.
124/// This function must be called in a build script.
125///
126/// The directory path must be a subdirectory of the manifest directory. The name of the bundle
127/// must later be used as an argument to the `include_fs!` macro.
128///
129/// # Example
130///
131/// ```rust
132/// // In build.rs
133/// include_fs::bundle("assets", "assets").unwrap();
134/// include_fs::bundle("./static/public", "public").unwrap();
135///
136/// // In main.rs
137/// static ASSETS: IncludeFs = include_fs!("assets");
138/// static PUBLIC: IncludeFs = include_fs!("public");
139/// ```
140pub fn bundle<P: AsRef<Path>>(dir: P, bundle_name: &str) -> Result<(), ArchiveError> {
141  let manifest_dir = env::var("CARGO_MANIFEST_DIR").expect("no CARGO_MANIFEST_DIR");
142  let source_dir = Path::new(&manifest_dir).join(dir).canonicalize()?;
143
144  // Ensure the source directory is a subdirectory of the manifest directory
145  if !source_dir.starts_with(&manifest_dir) {
146    return Err(ArchiveError::InvalidSourceDirectory);
147  }
148
149  let relative_source_dir = source_dir.strip_prefix(&manifest_dir).unwrap();
150  println!("cargo:rerun-if-changed={}", relative_source_dir.display());
151
152  let mut files = Vec::new();
153  let walk = WalkDir::new(&source_dir).follow_links(false);
154  for entry in walk {
155    let entry = entry?;
156    let meta = entry.metadata()?;
157    if !meta.is_file() {
158      continue;
159    }
160
161    let path = entry.path().strip_prefix(&manifest_dir).unwrap();
162    files.push(FileEntry::new(path, meta.len()));
163  }
164
165  let out_dir = env::var("OUT_DIR").expect("no OUT_DIR");
166  let output_file = format!("{}.embed_fs", bundle_name);
167  let output_path = Path::new(&out_dir).join(output_file);
168
169  write_archive(&files, &output_path)
170}
171
172struct FsEntry {
173  size: u64,
174  data_offset: u64,
175}
176
177impl FsEntry {
178  fn new(size: u64, data_offset: u64) -> Self {
179    Self { size, data_offset }
180  }
181}
182
183/// A lazy-loaded file system embedded in the binary.
184///
185/// The index will be parsed the first time it is accessed. Since only filenames are read on initialization, this should be very fast.
186/// To make sure the index is not read in a time-critical path, the lock can be manually initialized beforehand:
187///
188/// ```rust
189/// static ASSETS: IncludeFs = include_fs!("assets");
190///
191/// // This will block until the index fs is loaded.
192/// let _ = &*ASSETS;
193/// ```
194pub type IncludeFs = LazyLock<IncludeFsInner>;
195
196pub struct IncludeFsInner {
197  file_index: HashMap<String, FsEntry>,
198  archive_bytes: &'static [u8],
199}
200
201impl IncludeFsInner {
202  /// Initialize a new IncludeFs from the given bytes.
203  ///
204  /// This function is only meant to be called by the `include_fs!` macro.
205  pub fn new(archive_bytes: &'static [u8]) -> Result<Self, FsError> {
206    if &archive_bytes[0..4] != MAGIC {
207      return Err(FsError::InvalidArchive);
208    }
209
210    let file_count = u32::from_le_bytes([
211      archive_bytes[4],
212      archive_bytes[5],
213      archive_bytes[6],
214      archive_bytes[7],
215    ]) as usize;
216
217    let mut offset = 8;
218    let mut file_index = HashMap::with_capacity(file_count);
219
220    for _ in 0..file_count {
221      let path_len =
222        u16::from_le_bytes([archive_bytes[offset], archive_bytes[offset + 1]]) as usize;
223      offset += 2;
224
225      let path = String::from_utf8_lossy(&archive_bytes[offset..offset + path_len]).to_string();
226      offset += path_len;
227
228      let size = u64::from_le_bytes([
229        archive_bytes[offset],
230        archive_bytes[offset + 1],
231        archive_bytes[offset + 2],
232        archive_bytes[offset + 3],
233        archive_bytes[offset + 4],
234        archive_bytes[offset + 5],
235        archive_bytes[offset + 6],
236        archive_bytes[offset + 7],
237      ]);
238      offset += 8;
239
240      let data_offset = u64::from_le_bytes([
241        archive_bytes[offset],
242        archive_bytes[offset + 1],
243        archive_bytes[offset + 2],
244        archive_bytes[offset + 3],
245        archive_bytes[offset + 4],
246        archive_bytes[offset + 5],
247        archive_bytes[offset + 6],
248        archive_bytes[offset + 7],
249      ]);
250      offset += 8;
251
252      file_index.insert(path, FsEntry::new(size, data_offset));
253    }
254
255    Ok(IncludeFsInner {
256      file_index,
257      archive_bytes,
258    })
259  }
260
261  pub fn exists(&self, path: &str) -> bool {
262    self.file_index.contains_key(path)
263  }
264
265  pub fn get(&self, path: &str) -> Result<&[u8], FsError> {
266    let Some(entry) = self.file_index.get(path) else {
267      return Err(FsError::NotFound);
268    };
269
270    let start = entry.data_offset as usize;
271    let end = start + entry.size as usize;
272    Ok(&self.archive_bytes[start..end])
273  }
274
275  pub fn list_paths(&self) -> Vec<&str> {
276    self.file_index.keys().map(|s| s.as_str()).collect()
277  }
278}
279
280#[cfg(test)]
281mod tests {
282  use super::*;
283
284  #[test]
285  fn test_compute_header() {
286    let files = vec![
287      FileEntry::new("src/main.rs", 1024),
288      FileEntry::new("assets/image.png", 2048),
289    ];
290
291    let header = compute_header(&files).unwrap();
292
293    // Verify magic
294    assert_eq!(&header[0..4], b"INFS");
295
296    // Verify file count
297    let file_count = u32::from_le_bytes([header[4], header[5], header[6], header[7]]);
298    assert_eq!(file_count, 2);
299
300    // Basic size check (exact calculation depends on path lengths)
301    let expected_min_size = 4 + 4 + // magic + count
302      2 + "src/main.rs".len() + 8 + 8 + // first file
303      2 + "assets/image.png".len() + 8 + 8; // second file
304
305    assert_eq!(header.len(), expected_min_size);
306  }
307
308  #[test]
309  fn test_path_too_long() {
310    let long_path = "a".repeat(u16::MAX as usize + 1);
311    let files = vec![FileEntry::new(long_path.clone(), 100)];
312
313    let result = compute_header(&files);
314    assert!(matches!(result, Err(ArchiveError::PathTooLong { .. })));
315
316    if let Err(ArchiveError::PathTooLong { path, len, max }) = result {
317      assert_eq!(path, long_path);
318      assert_eq!(len, u16::MAX as usize + 1);
319      assert_eq!(max, u16::MAX as usize);
320    }
321  }
322}
include_fs/lib.rs

include_fs/
lib.rs
