diff options
Diffstat (limited to 'src/write.rs')
| -rw-r--r-- | src/write.rs | 368 |
1 files changed, 368 insertions, 0 deletions
diff --git a/src/write.rs b/src/write.rs new file mode 100644 index 0000000..23579b2 --- /dev/null +++ b/src/write.rs @@ -0,0 +1,368 @@ +use std::cell::RefCell; +use std::collections::{BTreeMap, HashMap}; +use std::ffi::{CStr, CString, OsStr, OsString}; +use std::io::Read; +use std::path::Path; +use std::time::SystemTime; +use bindings::*; +use super::*; +use super::SquashfsError; +use thiserror::Error; +use walkdir::{DirEntry, WalkDir}; + +pub mod BlockFlags { + pub const DontCompress: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_COMPRESS; + pub const BlockAlign: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_ALIGN; + pub const DontFragment: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_FRAGMENT; + pub const DontDeduplicate: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_DEDUPLICATE; + pub const IgnoreSparse: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_IGNORE_SPARSE; + pub const DontHash: u32 = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_HASH; +} + +pub enum SourceData { + File(Box<dyn Read>), + Dir(Box<dyn Iterator<Item=(OsString, u32)>>), + Symlink(OsString), + BlockDev(u32, u32), + CharDev(u32, u32), + Fifo, + Socket, +} + +pub struct Source { + data: SourceData, + xattrs: HashMap<OsString, Vec<u8>>, + uid: u32, + gid: u32, + mode: u16, + modified: u32, + flags: u32, +} + +fn copy_metadata(src: &ManagedPointer<sqfs_inode_generic_t>, dst: &mut ManagedPointer<sqfs_inode_generic_t>) { + fn nlink_ref(inode: &ManagedPointer<sqfs_inode_generic_t>) -> Option<&u32> { + unimplemented!(); + } + let (src_base, dst_base) = unsafe { (&(***src).base, &mut (***dst).base) }; + dst_base.mode = src_base.mode; + dst_base.uid_idx = src_base.uid_idx; + dst_base.gid_idx = src_base.gid_idx; + dst_base.mod_time = src_base.mod_time; + dst_base.inode_number = src_base.inode_number; + // TODO xattr_idx, uid, git, mode, mtime, link_count +} + +impl Source { + pub fn new(data: SourceData, xattrs: HashMap<OsString, Vec<u8>>, parent: u32, flags: u32) -> Self { // TODO Parent not necessary? + Self { data: data, xattrs: xattrs, uid: 1000, gid: 1001, mode: 0x1ff, modified: 0, flags: flags } + } + + fn devno(maj: u32, min: u32) -> u32 { + ((min & 0xfff00) << 20) | ((maj & 0xfff) << 8) | (min & 0xff) + } + + // TODO Handle hard links + fn to_inode(&self, link_count: u32) -> Result<ManagedPointer<sqfs_inode_generic_t>> { + fn create_inode(kind: SQFS_INODE_TYPE, extra: usize) -> ManagedPointer<sqfs_inode_generic_t> { + use std::alloc::{alloc, Layout}; + use std::mem::{align_of, size_of}; + unsafe { + let layout = Layout::from_size_align_unchecked(size_of::<sqfs_inode_generic_t>() + extra, align_of::<sqfs_inode_generic_t>()); + let ret = alloc(layout) as *mut sqfs_inode_generic_t; + (*ret).base.type_ = kind as u16; + ManagedPointer::new(ret, rust_dealloc) + } + } + let ret = unsafe { + match &self.data { + SourceData::File(_) => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_FILE, 0); + ret + }, + SourceData::Dir(_) => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_DIR, 0); + (**ret).data.dir.nlink = link_count; + ret + }, + SourceData::Symlink(dest_os) => { + let dest = os_to_string(&dest_os)?.into_bytes(); + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_SLINK, dest.len()); + let mut data = &mut (**ret).data.slink; + data.nlink = link_count; + data.target_size = dest.len() as u32; + let dest_field = std::mem::transmute::<_, &mut [u8]>((**ret).extra.as_mut_slice(dest.len())); + dest_field.copy_from_slice(dest.as_slice()); + ret + }, + SourceData::BlockDev(maj, min) => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_BDEV, 0); + let mut data = &mut (**ret).data.dev; + data.nlink = link_count; + data.devno = Self::devno(*maj, *min); + ret + }, + SourceData::CharDev(maj, min) => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_CDEV, 0); + let mut data = &mut (**ret).data.dev; + data.nlink = link_count; + data.devno = Self::devno(*maj, *min); + ret + }, + SourceData::Fifo => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_FIFO, 0); + (**ret).data.ipc.nlink = link_count; + ret + }, + SourceData::Socket => { + let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_SOCKET, 0); + (**ret).data.ipc.nlink = link_count; + ret + }, + } + }; + Ok(ret) + } +} + +struct IntermediateNode { + inode: Box<ManagedPointer<sqfs_inode_generic_t>>, + dir_children: Option<Box<dyn Iterator<Item=(OsString, u32)>>>, + pos: u64, + parent: u32, // TODO Calculate rather than requiring +} + +pub struct Writer { + outfile: ManagedPointer<sqfs_file_t>, + compressor_config: sqfs_compressor_config_t, + compressor: ManagedPointer<sqfs_compressor_t>, + superblock: sqfs_super_t, + block_writer: ManagedPointer<sqfs_block_writer_t>, + block_processor: ManagedPointer<sqfs_block_processor_t>, + frag_table: ManagedPointer<sqfs_frag_table_t>, + id_table: ManagedPointer<sqfs_id_table_t>, + xattr_writer: ManagedPointer<sqfs_xattr_writer_t>, + inode_writer: ManagedPointer<sqfs_meta_writer_t>, + dirent_writer: ManagedPointer<sqfs_meta_writer_t>, + dir_writer: ManagedPointer<sqfs_dir_writer_t>, + nodes: Vec<RefCell<IntermediateNode>>, +} + +impl Writer { + pub fn open(path: &Path) -> Result<Self> { + let cpath = CString::new(os_to_string(path.as_os_str())?)?; + let block_size = SQFS_DEFAULT_BLOCK_SIZE as u64; + let num_workers = 1; + let compressor_id = SQFS_COMPRESSOR_SQFS_COMP_ZSTD; + let now = 0; // TODO Get current timestamp + let outfile = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_open_file(cpath.as_ptr(), SQFS_FILE_OPEN_FLAGS_SQFS_FILE_OPEN_OVERWRITE) + }, &format!("Couldn't open output file {}", path.display()))?, sfs_destroy); + let compressor_config = sfs_init(&|x| unsafe { + sqfs_compressor_config_init(x, compressor_id, block_size, 0) + }, "Couldn't create compressor config")?; + let compressor = ManagedPointer::new(sfs_init_ptr(&|x| unsafe { + sqfs_compressor_create(&compressor_config, x) + }, "Couldn't create compressor")?, sfs_destroy); + let superblock = sfs_init(&|x| unsafe { + sqfs_super_init(x, block_size, now, compressor_id) + }, "Couldn't create superblock")?; + let frag_table = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_frag_table_create(0) + }, "Couldn't create fragment table")?, sfs_destroy); + let block_writer = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_block_writer_create(*outfile, 4096, 0) + }, "Couldn't create block writer")?, sfs_destroy); + let block_processor = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_block_processor_create(block_size, *compressor, num_workers, 10 * num_workers as u64, *block_writer, *frag_table) + }, "Couldn't create block processor")?, sfs_destroy); + let id_table = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_id_table_create(0) + }, "Couldn't create ID table")?, sfs_destroy); + let xattr_writer = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_xattr_writer_create(0) + }, "Couldn't create xattr writer")?, sfs_destroy); + let inode_writer = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_meta_writer_create(*outfile, *compressor, 0) + }, "Couldn't create inode metadata writer")?, sfs_destroy); + let dirent_writer = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_meta_writer_create(*outfile, *compressor, SQFS_META_WRITER_FLAGS_SQFS_META_WRITER_KEEP_IN_MEMORY) // TODO Untangle so we don't have to keep in memory + }, "Couldn't create directory entry metadata writer")?, sfs_destroy); + let dir_writer = ManagedPointer::new(sfs_init_check_null(&|| unsafe { + sqfs_dir_writer_create(*dirent_writer, SQFS_DIR_WRITER_CREATE_FLAGS_SQFS_DIR_WRITER_CREATE_EXPORT_TABLE) + }, "Couldn't create directory writer")?, sfs_destroy); + unsafe { + sfs_check(sqfs_super_write(&superblock, *outfile), "Couldn't write archive superblock")?; + sfs_check((**compressor).write_options.expect("Compressor doesn't provide write_options")(*compressor, *outfile), "Couldn't write compressor options")?; + } + Ok(Self { + outfile: outfile, + compressor_config: compressor_config, + compressor: compressor, + superblock: superblock, + block_writer: block_writer, + block_processor: block_processor, + frag_table: frag_table, + id_table: id_table, + xattr_writer: xattr_writer, + inode_writer: inode_writer, + dirent_writer: dirent_writer, + dir_writer: dir_writer, + nodes: vec![], + }) + } + + fn mode_from_inode(inode: &ManagedPointer<sqfs_inode_generic_t>) -> u16 { + let typenums = vec![ // TODO Lazy static + (SQFS_INODE_TYPE_SQFS_INODE_DIR, S_IFDIR), + (SQFS_INODE_TYPE_SQFS_INODE_FILE, S_IFREG), + (SQFS_INODE_TYPE_SQFS_INODE_SLINK, S_IFLNK), + (SQFS_INODE_TYPE_SQFS_INODE_BDEV, S_IFBLK), + (SQFS_INODE_TYPE_SQFS_INODE_CDEV, S_IFCHR), + (SQFS_INODE_TYPE_SQFS_INODE_FIFO, S_IFIFO), + (SQFS_INODE_TYPE_SQFS_INODE_SOCKET, S_IFSOCK), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_DIR, S_IFDIR), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_FILE, S_IFREG), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_SLINK, S_IFLNK), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_BDEV, S_IFBLK), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_CDEV, S_IFCHR), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_FIFO, S_IFIFO), + (SQFS_INODE_TYPE_SQFS_INODE_EXT_SOCKET, S_IFSOCK), + ].into_iter().collect::<HashMap<u32, u32>>(); + let base = unsafe { (***inode).base }; + typenums[&(base.type_ as u32)] as u16 | base.mode + } + + fn outfile_size(&self) -> u64 { + unsafe { (**self.outfile).get_size.expect("Superblock doesn't provide get_size")(*self.outfile) } + } + + // TODO Minimize unsafe blocks + pub fn add(&mut self, mut source: Source) -> Result<u32> { + let flags = source.flags; + let nlink = 1; // TODO Handle hard links + let mut inode = match source.data { + SourceData::File(ref mut reader) => { + let mut ret = Box::new(ManagedPointer::null(libc_free)); + unsafe { + sfs_check(sqfs_block_processor_begin_file(*self.block_processor, &mut **ret, ptr::null_mut(), flags), "Couldn't begin writing file")?; + let mut buf = vec![0; BLOCK_BUF_SIZE]; + loop { + let rdsize = reader.read(&mut buf)? as u64; + if rdsize == 0 { break; } + sfs_check(sqfs_block_processor_append(*self.block_processor, &buf as &[u8] as *const [u8] as *const libc::c_void, rdsize), "Couldn't write file data block")?; + } + sfs_check(sqfs_block_processor_end_file(*self.block_processor), "Couldn't finish writing file")?; + } + ret + }, + _ => Box::new(source.to_inode(nlink)?), + }; + unsafe { + sfs_check(sqfs_xattr_writer_begin(*self.xattr_writer, 0), "Couldn't start writing xattrs")?; + for (key, value) in &source.xattrs { + let ckey = CString::new(os_to_string(key)?)?; + sfs_check(sqfs_xattr_writer_add(*self.xattr_writer, ckey.as_ptr() as *const i8, value as &[u8] as *const [u8] as *const libc::c_void, value.len() as u64), "Couldn't add xattr")?; + } + let xattr_idx = unsafe { sfs_init(&|x| sqfs_xattr_writer_end(*self.xattr_writer, x), "Couldn't finish writing xattrs")? }; + let mut base = &mut (***inode).base; + base.mode = source.mode; + sqfs_inode_set_xattr_index(**inode, xattr_idx); + sfs_check(sqfs_id_table_id_to_index(*self.id_table, source.uid, &mut base.uid_idx), "Couldn't set inode UID"); + sfs_check(sqfs_id_table_id_to_index(*self.id_table, source.gid, &mut base.gid_idx), "Couldn't set inode GID"); + base.mod_time = source.modified; + base.inode_number = self.nodes.len() as u32 + 1;; + } + let dir_children = match source.data { + SourceData::Dir(children) => Some(children), + _ => None, + }; + self.nodes.push(RefCell::new(IntermediateNode { inode: inode, dir_children: dir_children, pos: 0, parent: 0 })); + Ok(self.nodes.len() as u32) + } + + pub fn finish(&mut self) -> Result<()> { + unsafe { + sfs_check(sqfs_block_processor_finish(*self.block_processor), "Failed finishing block processing")?; + self.superblock.inode_table_start = self.outfile_size(); + for raw_node in &self.nodes { + let mut node = raw_node.borrow_mut(); + // TODO Handle extended inodes properly + // TODO What happens if a dir tries to include itself as a child? Probably a RefCell borrow panic. + let id = (***node.inode).base.inode_number; + if let Some(children) = node.dir_children.take() { + sfs_check(sqfs_dir_writer_begin(*self.dir_writer, 0), "Couldn't start writing directory")?; + // For each child, need: name, ID, reference, mode + for (name, child_id) in children { // TODO Check that children are sorted + if child_id >= id { panic!("Tried to write directory {} before child {}", id, child_id) } // TODO Allocate error + let child_node = &self.nodes[child_id as usize - 1].borrow(); + let child = child_node.inode.as_ref(); + let child_ref = child_node.pos; + sfs_check(sqfs_dir_writer_add_entry(*self.dir_writer, CString::new(os_to_string(&name)?)?.as_ptr(), child_id, child_ref, Self::mode_from_inode(&child)), "Couldn't add directory entry")?; + } + sfs_check(sqfs_dir_writer_end(*self.dir_writer), "Couldn't finish writing directory")?; + let mut ret = Box::new(ManagedPointer::new(sfs_init_check_null(&|| sqfs_dir_writer_create_inode(*self.dir_writer, 0, 0, node.parent), "Couldn't get inode for directory")?, libc_free)); + copy_metadata(&*node.inode, &mut ret); + node.inode = ret; + } + let (mut block, mut offset) = (0, 0); + sqfs_meta_writer_get_position(*self.inode_writer, &mut block, &mut offset); + node.pos = block << 16 | offset as u64; + sfs_check(sqfs_meta_writer_write_inode(*self.inode_writer, **node.inode), "Couldn't write inode")?; + } + + let root_ref = self.nodes.last().ok_or(SquashfsError::Empty)?.borrow().pos; + self.superblock.root_inode_ref = root_ref; + sfs_check(sqfs_meta_writer_flush(*self.inode_writer), "Couldn't flush inodes")?; + sfs_check(sqfs_meta_writer_flush(*self.dirent_writer), "Couldn't flush directory entries")?; + self.superblock.directory_table_start = self.outfile_size(); + sfs_check(sqfs_meta_write_write_to_file(*self.dirent_writer), "Couldn't write directory entries")?; + (self.superblock).inode_count = self.nodes.len() as u32; + sfs_check(sqfs_frag_table_write(*self.frag_table, *self.outfile, &mut self.superblock, *self.compressor), "Couldn't write fragment table")?; + sfs_check(sqfs_dir_writer_write_export_table(*self.dir_writer, *self.outfile, *self.compressor, self.nodes.len() as u32, root_ref, &mut self.superblock), "Couldn't write export table")?; + sfs_check(sqfs_id_table_write(*self.id_table, *self.outfile, &mut self.superblock, *self.compressor), "Couldn't write ID table")?; + sfs_check(sqfs_xattr_writer_flush(*self.xattr_writer, *self.outfile, &mut self.superblock, *self.compressor), "Couldn't write xattr table")?; + self.superblock.bytes_used = self.outfile_size(); + self.superblock.modification_time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_secs() as u32; + sfs_check(sqfs_super_write(&self.superblock, *self.outfile), "Couldn't rewrite archive superblock")?; + let padding: Vec<u8> = vec![0; PAD_TO - self.outfile_size() as usize % PAD_TO]; + sfs_check((**self.outfile).write_at.expect("File does not provide write_at")(*self.outfile, self.outfile_size(), &padding as &[u8] as *const [u8] as *const libc::c_void, padding.len() as u64), "Couldn't pad file"); + } + Ok(()) + } + + pub fn add_tree<T: AsRef<Path>>(&mut self, root: T, callback: &Fn(Source) -> Result<Source>) -> Result<()> { + let mut childmap: HashMap<PathBuf, BTreeMap<OsString, u32>> = HashMap::new(); + for step in WalkDir::new(root.as_ref()).follow_links(false).contents_first(true) { + match step { + Ok(entry) => { + // TODO Consider adding Unix-specific functionality with graceful degradation + // TODO Catch all errors except add() and continue + let metadata = entry.metadata().unwrap(); + let mtime = metadata.modified()?.duration_since(SystemTime::UNIX_EPOCH)?.as_secs() as u32; + let data = if metadata.file_type().is_dir() { + SourceData::Dir(Box::new(childmap.remove(&entry.path().to_path_buf()).unwrap_or(BTreeMap::new()).into_iter())) + } + else if metadata.file_type().is_file() { + SourceData::File(Box::new(std::fs::File::open(entry.path())?)) + } + else if metadata.file_type().is_symlink() { + SourceData::Symlink(std::fs::read_link(entry.path())?.into_os_string()) + } + else { + panic!("Unknown or unsupported file type"); // TODO Error + }; + let id = self.add(callback(Source { data: data, xattrs: HashMap::new(), uid: 0, gid: 0, mode: 0x1ff, modified: mtime, flags: 0 })?)?; + if let Some(parent) = entry.path().parent() { + childmap.entry(parent.to_path_buf()).or_insert(BTreeMap::new()).insert(entry.file_name().to_os_string(), id); + } + println!("{}: {}", id, entry.path().display()); + }, + Err(e) => { + let path = e.path().map(|x| x.to_string_lossy().into_owned()).unwrap_or("(unknown)".to_string()); + eprintln!("Not processing {}: {}", path, e.to_string()); + } + } + } + Ok(()) + } +} |