use std::cell::RefCell;
use std::collections::{BTreeMap, HashMap};
use std::ffi::CString;
use std::fs::Metadata;
use std::io::Read;
use std::path::Path;
use std::sync::{Mutex, RwLock};
use std::time::SystemTime;
use super::*;
use walkdir::{DirEntry, WalkDir};
#[repr(u32)]
pub enum BlockFlags {
DontCompress = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_COMPRESS,
BlockAlign = super::SQFS_BLK_FLAGS_SQFS_BLK_ALIGN,
DontFragment = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_FRAGMENT,
DontDeduplicate = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_DEDUPLICATE,
IgnoreSparse = super::SQFS_BLK_FLAGS_SQFS_BLK_IGNORE_SPARSE,
DontHash = super::SQFS_BLK_FLAGS_SQFS_BLK_DONT_HASH,
}
pub enum SourceData {
File(Box<dyn Read + Sync + Send>),
Dir(Box<dyn Iterator<Item=(OsString, u32)> + Sync + Send>),
Symlink(PathBuf),
BlockDev(u32, u32),
CharDev(u32, u32),
Fifo,
Socket,
}
pub struct Source {
pub data: SourceData,
pub uid: u32,
pub gid: u32,
pub mode: u16,
pub modified: u32,
pub xattrs: HashMap<OsString, Vec<u8>>,
pub flags: u32,
}
fn file_xattrs(path: &Path) -> Result<HashMap<OsString, Vec<u8>>> {
xattr::list(path)?.map(|attr| {
let value = xattr::get(path, attr.clone()).map_err(|e| SquashfsError::Xattr(path.to_path_buf(), e))?
.expect(&format!("Could not retrieve xattr {:?} reported to be present", attr));
Ok((attr, value))
}).collect()
}
fn copy_metadata(src: &ManagedPointer<sqfs_inode_generic_t>, dst: &mut ManagedPointer<sqfs_inode_generic_t>) -> Result<()> {
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;
let mut xattr_idx: u32 = 0;
unsafe {
sfs_check(sqfs_inode_get_xattr_index(**src, &mut xattr_idx), "Couldn't get xattr index")?;
sfs_check(sqfs_inode_set_xattr_index(**dst, xattr_idx), "Couldn't set xattr index")?;
}
Ok(())
}
impl Source {
pub fn defaults(data: SourceData) -> Self {
Self { data: data, uid: 0, gid: 0, mode: 0x1ff, modified: 0, xattrs: HashMap::new(), flags: 0 }
}
fn devno(maj: u32, min: u32) -> u32 {
((min & 0xfff00) << 20) | ((maj & 0xfff) << 8) | (min & 0xff)
}
unsafe fn to_inode(&self, link_count: u32) -> Result<ManagedPointer<sqfs_inode_generic_t>> { unsafe {
unsafe fn create_inode(kind: SQFS_INODE_TYPE, extra: usize) -> ManagedPointer<sqfs_inode_generic_t> { unsafe {
use std::alloc::{alloc, Layout};
use std::mem::{align_of, size_of};
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 = match &self.data {
SourceData::File(_) => create_inode(SQFS_INODE_TYPE_SQFS_INODE_FILE, 0),
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.as_os_str())?.into_bytes();
let mut ret = create_inode(SQFS_INODE_TYPE_SQFS_INODE_SLINK, dest.len());
let 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 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 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)> + Sync + Send>>,
pos: u64,
}
pub struct SourceFile {
pub path: PathBuf,
pub content: Source,
}
impl SourceFile {
pub fn defaults(path: PathBuf, data: SourceData) -> Self {
Self { path: path, content: Source::defaults(data) }
}
}
pub struct Writer {
outfile: ManagedPointer<sqfs_file_t>,
#[allow(dead_code)] compressor_config: sqfs_compressor_config_t, compressor: ManagedPointer<sqfs_compressor_t>,
superblock: sqfs_super_t,
#[allow(dead_code)] block_writer: ManagedPointer<sqfs_block_writer_t>, block_processor: Mutex<ManagedPointer<sqfs_block_processor_t>>,
frag_table: ManagedPointer<sqfs_frag_table_t>,
id_table: Mutex<ManagedPointer<sqfs_id_table_t>>,
xattr_writer: Mutex<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: Mutex<Vec<RefCell<IntermediateNode>>>,
finished: RwLock<bool>,
}
impl Writer {
pub fn open<T: AsRef<Path>>(path: T) -> Result<Self> {
let cpath = CString::new(os_to_string(path.as_ref().as_os_str())?)?;
let block_size = SQFS_DEFAULT_BLOCK_SIZE as usize;
let num_workers = num_cpus::get() as u32;
let compressor_id = SQFS_COMPRESSOR_SQFS_COMP_ZSTD;
let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_secs() as u32;
let outfile = 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.as_ref().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 = 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 = sfs_init_check_null(&|| unsafe {
sqfs_frag_table_create(0)
}, "Couldn't create fragment table", sfs_destroy)?;
let block_writer = sfs_init_check_null(&|| unsafe {
sqfs_block_writer_create(*outfile, 4096, 0)
}, "Couldn't create block writer", sfs_destroy)?;
let block_processor = Mutex::new(sfs_init_check_null(&|| unsafe {
sqfs_block_processor_create(block_size, *compressor, num_workers, 10 * num_workers as usize, *block_writer, *frag_table)
}, "Couldn't create block processor", sfs_destroy)?);
let id_table = Mutex::new(sfs_init_check_null(&|| unsafe {
sqfs_id_table_create(0)
}, "Couldn't create ID table", sfs_destroy)?);
let xattr_writer = Mutex::new(sfs_init_check_null(&|| unsafe {
sqfs_xattr_writer_create(0)
}, "Couldn't create xattr writer", sfs_destroy)?);
let inode_writer = sfs_init_check_null(&|| unsafe {
sqfs_meta_writer_create(*outfile, *compressor, 0)
}, "Couldn't create inode metadata writer", sfs_destroy)?;
let dirent_writer = sfs_init_check_null(&|| unsafe {
sqfs_meta_writer_create(*outfile, *compressor, SQFS_META_WRITER_FLAGS_SQFS_META_WRITER_KEEP_IN_MEMORY)
}, "Couldn't create directory entry metadata writer", sfs_destroy)?;
let dir_writer = 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: Mutex::new(vec![]),
finished: RwLock::new(false),
})
}
fn mode_from_inode(inode: &ManagedPointer<sqfs_inode_generic_t>) -> u16 {
lazy_static! {
static ref TYPENUMS: HashMap<u32, u32> = vec![
(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();
}
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) }
}
pub fn add(&mut self, mut source: Source) -> Result<u32> {
let finished = self.finished.read().expect("Poisoned lock");
if *finished { Err(SquashfsError::Finished)?; }
let flags = source.flags;
let nlink = 1; let mut inode = unsafe {
match source.data {
SourceData::File(ref mut reader) => {
let mut ret = Box::new(ManagedPointer::null(libc_free));
let block_processor = self.block_processor.lock().expect("Poisoned lock");
sfs_check(sqfs_block_processor_begin_file(**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)?;
if rdsize == 0 { break; }
sfs_check(sqfs_block_processor_append(**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(**block_processor), "Couldn't finish writing file")?;
ret
},
_ => Box::new(source.to_inode(nlink)?),
}
};
unsafe {
let xattr_writer = self.xattr_writer.lock().expect("Poisoned lock");
sfs_check(sqfs_xattr_writer_begin(**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(**xattr_writer, ckey.as_ptr() as *const libc::c_char, value.as_ptr() as *const libc::c_void, value.len()), "Couldn't add xattr")?;
}
let xattr_idx = sfs_init(&|x| sqfs_xattr_writer_end(**xattr_writer, x), "Couldn't finish writing xattrs")?;
let base = &mut (***inode).base;
base.mode = source.mode;
sqfs_inode_set_xattr_index(**inode, xattr_idx);
let id_table = self.id_table.lock().expect("Poisoned lock");
sfs_check(sqfs_id_table_id_to_index(**id_table, source.uid, &mut base.uid_idx), "Couldn't set inode UID")?;
sfs_check(sqfs_id_table_id_to_index(**id_table, source.gid, &mut base.gid_idx), "Couldn't set inode GID")?;
base.mod_time = source.modified;
}
let dir_children = match source.data {
SourceData::Dir(children) => Some(children),
_ => None,
};
let mut nodes = self.nodes.lock().expect("Poisoned lock");
let nodenum = nodes.len() as u32 + 1;
unsafe { (***inode).base.inode_number = nodenum; }
nodes.push(RefCell::new(IntermediateNode { inode: inode, dir_children: dir_children, pos: 0 }));
Ok(nodenum)
}
pub fn finish(&mut self) -> Result<()> {
*self.finished.write().expect("Poisoned lock") = true;
let nodes = self.nodes.lock().expect("Poisoned lock");
unsafe {
sfs_check(sqfs_block_processor_finish(**self.block_processor.lock().expect("Poisoned lock")), "Failed finishing block processing")?;
self.superblock.inode_table_start = self.outfile_size();
for raw_node in &*nodes {
let mut node = raw_node.borrow_mut();
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 (name, child_id) in children { if child_id >= id { Err(SquashfsError::WriteOrder(child_id))?; }
let child_node = &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(sfs_init_check_null(&|| {
sqfs_dir_writer_create_inode(*self.dir_writer, 0, 0, 0) }, "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 = 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 = 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, nodes.len() as u32, root_ref, &mut self.superblock), "Couldn't write export table")?;
sfs_check(sqfs_id_table_write(**self.id_table.lock().expect("Poisoned lock"), *self.outfile, &mut self.superblock, *self.compressor), "Couldn't write ID table")?;
sfs_check(sqfs_xattr_writer_flush(**self.xattr_writer.lock().expect("Poisoned lock"), *self.outfile, &mut self.superblock, *self.compressor), "Couldn't write xattr table")?;
self.superblock.bytes_used = self.outfile_size();
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()), "Couldn't pad file")?;
}
Ok(())
}
}
unsafe impl Sync for Writer { }
unsafe impl Send for Writer { }
enum ChildMapEntry {
Accumulating(BTreeMap<OsString, u32>),
Done,
}
impl ChildMapEntry {
fn new() -> Self {
Self::Accumulating(BTreeMap::new())
}
fn add(&mut self, name: OsString, id: u32) -> Result<()> {
match self {
Self::Done => Err(SquashfsError::WriteOrder(id))?,
Self::Accumulating(children) => {
children.insert(name, id);
Ok(())
},
}
}
fn finish(&mut self) -> Result<BTreeMap<OsString, u32>> {
match std::mem::replace(self, Self::Done) {
Self::Done => Err(SquashfsError::Internal("Tried to finish directory in tree processor multiple times".to_string()))?,
Self::Accumulating(children) => Ok(children),
}
}
}
pub struct TreeProcessor {
writer: Mutex<Writer>,
childmap: Mutex<HashMap<PathBuf, ChildMapEntry>>,
}
impl TreeProcessor {
pub fn new<P: AsRef<Path>>(outfile: P) -> Result<Self> {
let writer = Writer::open(outfile)?;
Ok(Self { writer: Mutex::new(writer), childmap: Mutex::new(HashMap::new()) })
}
pub fn add(&self, mut source: SourceFile) -> Result<u32> {
let mut childmap = self.childmap.lock().expect("Poisoned lock");
if let SourceData::Dir(old_children) = &mut source.content.data {
let mut children = childmap.entry(source.path.clone()).or_insert(ChildMapEntry::new()).finish()?;
children.extend(old_children);
source.content.data = SourceData::Dir(Box::new(children.into_iter()));
}
let id = self.writer.lock().expect("Poisoned lock").add(source.content)?;
if let Some(parent) = source.path.parent() {
childmap.entry(parent.to_path_buf()).or_insert(ChildMapEntry::new()).add(source.path.file_name().expect("Path from walkdir has no file name").to_os_string(), id)?;
}
Ok(id)
}
#[cfg(target_os = "linux")]
fn apply_metadata(source: &mut Source, metadata: Metadata) {
use std::os::linux::fs::MetadataExt;
source.uid = metadata.st_uid();
source.gid = metadata.st_gid();
source.mode = (metadata.st_mode() & !S_IFMT) as u16;
}
#[cfg(all(not(target_os = "linux"), target_family = "unix"))]
fn apply_metadata(source: &mut Source, metadata: Metadata) {
use std::os::unix::fs::MetadataExt;
source.uid = metadata.uid();
source.gid = metadata.gid();
source.mode = (metadata.mode() & 0x1ff) as u16;
}
#[cfg(all(not(target_os = "linux"), not(target_family = "unix")))]
fn apply_metadata(_source: &mut Source, _metadata: Metadata) { }
pub fn finish(&self) -> Result<()> {
self.writer.lock().expect("Poisoned lock").finish()
}
fn make_source(&self, entry: DirEntry) -> Result<Source> {
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(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())?)
}
else {
Err(SquashfsError::WriteType(metadata.file_type()))?;
unreachable!();
};
let mut source = Source { data: data, xattrs: file_xattrs(entry.path())?, uid: 0, gid: 0, mode: 0x1ff, modified: mtime, flags: 0 };
Self::apply_metadata(&mut source, metadata);
Ok(source)
}
pub fn iter<'a, P: AsRef<Path>>(&'a self, root: P) -> TreeIterator<'a> {
let tree = WalkDir::new(root).follow_links(false).contents_first(true);
TreeIterator { processor: self, tree: tree.into_iter() }
}
pub fn process<P: AsRef<Path>>(self, root: P) -> Result<()> {
for entry in self.iter(root) { self.add(entry?)?; }
self.finish()?;
Ok(())
}
}
pub struct TreeIterator<'a> {
processor: &'a TreeProcessor,
tree: walkdir::IntoIter,
}
impl<'a> std::iter::Iterator for TreeIterator<'a> {
type Item = Result<SourceFile>;
fn next(&mut self) -> Option<Self::Item> {
match self.tree.next() {
None => None,
Some(Ok(entry)) => {
let path = entry.path().to_path_buf();
Some(self.processor.make_source(entry).map(|source| SourceFile { path: path, content: source }))
},
Some(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());
self.next()
},
}
}
}