selfe_arc/
pack.rs

1use std::convert::TryFrom;
2use std::fs;
3use std::io::{self, Write};
4use std::path::{Path, PathBuf};
5use std::process::{Command, Stdio};
6
7use crate::layout;
8
9pub struct Archive {
10    files: Vec<File>,
11}
12
13pub struct File {
14    name: String,
15    path: PathBuf,
16}
17
18struct ScheduledFile {
19    path: PathBuf,
20    size: u64,
21    padding: u64,
22}
23
24#[derive(Debug)]
25pub enum ArchiveWriteError {
26    HeaderTooLarge,
27    DataSegmentTooLarge,
28    FileNameTooLong(String),
29    IO(io::Error),
30    UnsupportedTargetArch,
31    LinkError,
32}
33
34impl std::convert::From<io::Error> for ArchiveWriteError {
35    fn from(e: io::Error) -> ArchiveWriteError {
36        ArchiveWriteError::IO(e)
37    }
38}
39
40#[derive(Debug, Eq, PartialEq)]
41pub enum AddFileError {
42    EmptyNameNotAllowed,
43    NameConflict,
44    FileNameTooLong(String),
45}
46
47const LINKER_SCRIPT: &str = r#"SECTIONS
48{
49  .rodata : ALIGN(4096)
50  {
51    _selfe_arc_data_start = . ;
52    *(.*) ;
53    _selfe_arc_data_end = . ;
54  }
55}"#;
56
57impl Archive {
58    pub fn new() -> Archive {
59        Archive { files: vec![] }
60    }
61
62    pub fn add_file<P: AsRef<Path>>(&mut self, name: &str, path: P) -> Result<(), AddFileError> {
63        let path = path.as_ref();
64
65        if name.is_empty() {
66            return Err(AddFileError::EmptyNameNotAllowed);
67        }
68
69        if self.files.iter().find(|f| f.name == name).is_some() {
70            return Err(AddFileError::NameConflict);
71        }
72
73        if name.as_bytes().len() > layout::FILE_NAME_BYTES {
74            return Err(AddFileError::FileNameTooLong(name.to_owned()));
75        }
76
77        self.files.push(File {
78            name: name.to_owned(),
79            path: path.to_owned(),
80        });
81
82        Ok(())
83    }
84
85    pub fn write<W: Write>(&self, mut writer: &mut W) -> Result<(), ArchiveWriteError> {
86        let header_size = layout::ArchiveHeader::serialized_size();
87        let dir_entry_size = layout::DirectoryEntry::serialized_size();
88
89        let file_count =
90            u32::try_from(self.files.len()).map_err(|_| ArchiveWriteError::HeaderTooLarge)?;
91        let dir_size: u32 = file_count
92            .checked_mul(
93                u32::try_from(dir_entry_size).map_err(|_| ArchiveWriteError::HeaderTooLarge)?,
94            )
95            .ok_or(ArchiveWriteError::HeaderTooLarge)?;
96        let data_start: u32 = dir_size
97            .checked_add(u32::try_from(header_size).map_err(|_| ArchiveWriteError::HeaderTooLarge)?)
98            .ok_or(ArchiveWriteError::HeaderTooLarge)?;
99
100        // page align data_start
101        let data_start = layout::align_addr(data_start);
102        let initial_padding_size = data_start - (dir_size + header_size as u32);
103
104        // header
105        let header = layout::ArchiveHeader {
106            magic: *layout::MAGIC,
107            version: layout::VERSION_1,
108            data_start,
109            file_count,
110        };
111
112        header.write(&mut writer)?;
113
114        // directory
115        let mut scheduled_files = Vec::new();
116        let mut data_cursor = 0u64;
117        for (i, f) in self.files.iter().enumerate() {
118            // files should always be page-aligned
119            assert_eq!(data_cursor & 0xfff, 0);
120
121            let name = f.name.as_bytes();
122            if name.len() > layout::FILE_NAME_BYTES {
123                return Err(ArchiveWriteError::FileNameTooLong(f.name.to_owned()));
124            }
125
126            let data_file = fs::File::open(&f.path)?;
127            let file_size = data_file.metadata()?.len();
128
129            let mut dir_entry = layout::DirectoryEntry {
130                name_len: name.len() as u8,
131                name_bytes: [0; layout::FILE_NAME_BYTES],
132                offset: data_cursor,
133                length: file_size,
134            };
135
136            // copy the name into the dir entry
137            for (name_char, entry_char) in name.iter().zip(dir_entry.name_bytes.iter_mut()) {
138                *entry_char = *name_char;
139            }
140
141            dir_entry.write(&mut writer)?;
142
143            // pad to page boundaries, but not the last file.
144            let is_last = i == self.files.len() - 1;
145            let padding = if is_last {
146                0
147            } else {
148                let alignment: u64 = layout::ALIGNMENT.into();
149                let mask: u64 = layout::ALIGNMENT_MASK.into();
150                alignment - (file_size & mask)
151            };
152
153            scheduled_files.push(ScheduledFile {
154                path: f.path.to_owned(),
155                size: file_size,
156                padding,
157            });
158
159            data_cursor = data_cursor
160                .checked_add(dir_entry.length)
161                .ok_or(ArchiveWriteError::DataSegmentTooLarge)?
162                .checked_add(padding)
163                .ok_or(ArchiveWriteError::DataSegmentTooLarge)?;
164        }
165
166        // initial padding
167        for _ in 0..initial_padding_size {
168            writer.write_all(&[0])?;
169        }
170
171        // data
172        for f in scheduled_files.iter() {
173            let data_file = fs::File::open(&f.path).unwrap();
174            let mut buf_reader = io::BufReader::new(data_file);
175            let bytes_written = io::copy(&mut buf_reader, &mut writer)?;
176
177            assert_eq!(bytes_written, f.size, "Unexpected size for {:?}", f.path);
178
179            for _ in 0..f.padding {
180                writer.write_all(&[0])?;
181            }
182        }
183
184        Ok(())
185    }
186
187    pub fn write_object_file<P: AsRef<Path>, P2: AsRef<Path>>(
188        &self,
189        output: P,
190        ld: P2,
191        target_arch: &str,
192    ) -> Result<(), ArchiveWriteError> {
193        let output = output.as_ref();
194        let ld = ld.as_ref();
195
196        let archive_path = output.with_extension("selfearc");
197
198        {
199            let mut archive_file = fs::File::create(&archive_path)?;
200            self.write(&mut archive_file)?;
201        }
202
203        let linker_script_path = output.with_extension("ld");
204
205        {
206            let mut linker_script_file = fs::File::create(&*linker_script_path)?;
207            write!(&mut linker_script_file, "{}", LINKER_SCRIPT)?;
208        }
209
210        let output_format = match target_arch {
211            "aarch64" => "elf64-littleaarch64",
212            "arm" | "armv7" | "armebv7r" | "armv5te" | "armv7r" | "armv7s" => "elf32-littlearm",
213            "i386" | "i586" | "i686" => "elf32-i386",
214            "riscv32imac" | "riscv32imc" | "riscv64gc" | "riscv64imac" => "elf32-littleriscv",
215            "thumbv7em" | "thumbv7m" | "thumbv7neon" => "elf32-littlearm",
216            "thumbv8m.main" => "elf64-littleaarch64",
217            "x86_64" => "elf64-x86-64",
218            _ => return Err(ArchiveWriteError::UnsupportedTargetArch),
219        };
220
221        let mut ld = Command::new(ld);
222        ld.arg("-T")
223            .arg(linker_script_path)
224            .arg("--oformat")
225            .arg(output_format)
226            .arg("-r")
227            .arg("-b")
228            .arg("binary")
229            .arg(archive_path)
230            .arg("-o")
231            .arg(output)
232            .stdout(Stdio::inherit())
233            .stderr(Stdio::inherit());
234        println!("running ld command: {:?}", ld);
235
236        let output = ld.output()?;
237        if !output.status.success() {
238            return Err(ArchiveWriteError::LinkError);
239        }
240
241        Ok(())
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    #[test]
250    fn no_empty_name() {
251        let mut ar = Archive::new();
252        let res = ar.add_file("", Path::new("doesn't_matter"));
253        assert_eq!(res, Err(AddFileError::EmptyNameNotAllowed));
254    }
255
256    #[test]
257    fn no_duplicate_name() {
258        let dir = tempfile::tempdir().unwrap();
259        let f = dir.path().join("pack_test.txt");
260
261        {
262            let mut test_file = fs::File::create(&f).unwrap();
263            test_file.write_all(b"test").unwrap();
264            test_file.flush().unwrap();
265        }
266
267        let mut ar = Archive::new();
268        let res = ar.add_file("test", &f);
269        assert_eq!(res, Ok(()));
270
271        let res = ar.add_file("test", Path::new("doesn't_matter"));
272        assert_eq!(res, Err(AddFileError::NameConflict));
273    }
274
275    #[test]
276    fn no_overlong_name() {
277        let mut ar = Archive::new();
278        let name = "dajlsdkfj alskdjflkasdjfkljasdkl fjalfj eliwjf lasdijflaksdjflkasjdlkfaj sdlfkjasldkf jalsdkjf laskjdf laskdjf lakwjflawjelf ijasdlkfjaslfiawejlfajsdkflasdkjflaskdjflaskdjflaskdjflaksjdflkasjdflaksdjflaskdjflaksdjflkasjdflkajsdflkajsdlkfjasldkfjlaksjdflkasjdflkajsdlkfjasldkjfaklsdjf";
279        let res = ar.add_file(name, Path::new("foo"));
280        assert_eq!(res, Err(AddFileError::FileNameTooLong(name.to_owned())));
281    }
282
283    #[test]
284    fn pack_files() {
285        let dir = tempfile::tempdir().unwrap();
286        let f = dir.path().join("pack_test.txt");
287
288        {
289            let mut test_file = fs::File::create(&f).unwrap();
290            test_file.write_all(b"test").unwrap();
291            test_file.flush().unwrap();
292        }
293
294        let mut ar = Archive::new();
295        ar.add_file("test", &f).unwrap();
296
297        let mut actual_data = Vec::new();
298        {
299            let mut writer = io::BufWriter::new(&mut actual_data);
300            ar.write(&mut writer).unwrap();
301        }
302
303        let mut expected_data = vec![];
304        // ARCHIVE HEADER
305        #[rustfmt::skip]
306        expected_data.append(&mut vec!(
307            // magic
308            0x73, 0x65, 0x6c, 0x66, 0x65, 0x61, 0x72, 0x63,
309            // version
310            0x01,
311            // data_start
312            0x00, 0x10, 0x00, 0x00,
313            // file_count
314            0x01, 0x00, 0x00, 0x00,
315        ));
316
317        assert_eq!(
318            expected_data.len(),
319            layout::ArchiveHeader::serialized_size()
320        );
321
322        // DIRECTORY ENTRY 1/1
323        #[rustfmt::skip]
324        expected_data.append(&mut vec!(
325            // len, name
326            0x04, 0x74, 0x65, 0x73, 0x74, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
327            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
328            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
329            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
330            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
331            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
332            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
333            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
334            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
335            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
336            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
337            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
338            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
339            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
340            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
341            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
342            // offset
343            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
344            // length
345            0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
346        ));
347
348        assert_eq!(
349            expected_data.len(),
350            layout::ArchiveHeader::serialized_size() + layout::DirectoryEntry::serialized_size()
351        );
352
353        // PADDING
354        expected_data.append(&mut [0u8; 3807].to_vec());
355
356        // FILE 1/1
357        expected_data.append(&mut vec![0x74, 0x65, 0x73, 0x74]);
358
359        // double check file data alignment
360        assert_eq!(
361            expected_data
362                .clone()
363                .into_iter()
364                .skip(0x1000)
365                .collect::<Vec<u8>>(),
366            vec!(0x74, 0x65, 0x73, 0x74)
367        );
368
369        assert_eq!(expected_data.len(), actual_data.len());
370        for (i, (e, a)) in expected_data.iter().zip(actual_data.iter()).enumerate() {
371            assert_eq!(
372                e, a,
373                "At byte {:#x}, expected {:#04x} but got {:#04x}",
374                i, e, a
375            );
376        }
377    }
378
379    #[test]
380    fn object_file() {
381        use std::str;
382        let dir = tempfile::tempdir().unwrap();
383        let f = dir.path().join("pack_test.txt");
384        let elf = dir.path().join("pack_test.elf");
385
386        {
387            let mut test_file = fs::File::create(&f).unwrap();
388            test_file.write_all(b"test").unwrap();
389            test_file.flush().unwrap();
390        }
391
392        let mut ar = Archive::new();
393        ar.add_file("test", &f).unwrap();
394
395        ar.write_object_file(elf.to_str().unwrap(), "ld", "x86_64")
396            .unwrap();
397
398        let mut ld = Command::new("objdump");
399        let out = ld.arg("-t").arg(elf.to_str().unwrap()).output().unwrap();
400        let stdout = str::from_utf8(&out.stdout).unwrap();
401        assert!(stdout.contains("_selfe_arc_data_start"));
402        assert!(stdout.contains("_selfe_arc_data_end"));
403    }
404}
selfe_arc/pack.rs

selfe_arc/
pack.rs
