use std::collections::BTreeMap;
use std::ffi::{OsStr, OsString};
use std::io::{self, Read, Write};
use std::path::Path;
use std::process::ExitCode;
use std::sync::atomic::AtomicU64;
use std::sync::{Arc, Mutex};
use stuffr::zstd_opts::{self, CParamKey};
use super::CountedRun;
use super::host::{Host, SANDBOX_OUTPUT_CAP, Sandbox, reset_sigpipe};
pub(crate) mod cli;
mod dibio;
pub(crate) mod display;
pub(crate) mod fileio;
mod formats;
mod list;
mod messages;
pub(crate) mod util;
use cli::{Command, Limits, Operation, Progress, STDINMARK, STDOUTMARK, Stop};
use display::{Display, exm_throw};
use messages::{self as m, printf};
use util::Cores;
pub fn run(name: &'static str, argv0: &OsStr, args: Vec<OsString>) -> ExitCode {
reset_sigpipe();
let argv0 = if argv0.is_empty() {
name.as_bytes()
} else {
argv0.as_encoded_bytes()
};
let env = |k: &str| std::env::var_os(k);
let mut err = io::stderr();
let code = main(
argv0,
&args,
&env,
&mut Host::Real,
Cores::detect(),
None,
&mut err,
);
ExitCode::from(code as u8)
}
#[doc(hidden)]
pub fn run_in_counted(
dir: &Path,
name: &'static str,
args: &[OsString],
stdin: &[u8],
env: &dyn Fn(&str) -> Option<OsString>,
) -> CountedRun {
let stdout = Arc::new(Mutex::new(Vec::new()));
let reads = Arc::new(Mutex::new(BTreeMap::new()));
let mut host = Host::Sandbox(Sandbox {
dir: dir.to_path_buf(),
stdin: Some(stdin.to_vec()),
stdout: stdout.clone(),
budget: Arc::new(AtomicU64::new(SANDBOX_OUTPUT_CAP)),
reads: reads.clone(),
});
let mut err = Vec::new();
let code = main(
name.as_bytes(),
args,
env,
&mut host,
Cores::SANDBOX,
Some(Limits::SANDBOX),
&mut err,
);
drop(host);
let out = std::mem::take(&mut *stdout.lock().unwrap_or_else(|p| p.into_inner()));
let reads = std::mem::take(&mut *reads.lock().unwrap_or_else(|p| p.into_inner()));
(code, out, err, reads)
}
fn main(
argv0: &[u8],
args: &[OsString],
env: &dyn Fn(&str) -> Option<OsString>,
host: &mut Host,
cores: Cores,
limits: Option<Limits>,
err: &mut dyn Write,
) -> i32 {
let program = String::from_utf8_lossy(cli::last_name_from_path(argv0)).into_owned();
let mut out = host.out_writer(None);
let stop = match cli::parse(&program, env, args, cores, &mut out, err) {
Ok(mut cmd) => {
cmd.limits = limits;
execute(&mut cmd, host, &mut out, err)
}
Err(stop) => stop,
};
let _ = out.flush();
if stop.clean && stop.pause {
wait_enter(host, err);
}
stop.code
}
fn wait_enter(host: &mut Host, err: &mut dyn Write) {
printf(err, m::PRESS_ENTER, &[]);
let mut byte = [0u8; 1];
let _ = host.stdin().read(&mut byte);
}
fn is_link_not_fifo(host: &Host, name: &[u8]) -> bool {
util::is_link(host, name) && !util::is_fifo(host, name)
}
fn file_size(host: &Host, name: &[u8]) -> Option<u64> {
Some(util::file_size(host, name)).filter(|&s| s != util::FILESIZE_UNKNOWN)
}
fn finish(r: fileio::Fio<i32>, pause: bool) -> Stop {
match r {
Ok(code) => Stop {
code,
clean: true,
pause,
},
Err(fileio::Thrown(code)) => Stop {
code,
clean: false,
pause,
},
}
}
fn execute(cmd: &mut Command, host: &mut Host, out: &mut dyn Write, err: &mut dyn Write) -> Stop {
let pause = cmd.main_pause;
let clean = |code| Stop {
code,
clean: true,
pause,
};
let exit = |code| Stop {
code,
clean: false,
pause,
};
let mut d = Display::new(cmd.display_level);
cmd.prefs.util_level = cmd.display_level;
let util_level = cmd.display_level;
if !cmd.follow_links {
let total = cmd.files.len();
let mut kept = Vec::with_capacity(total);
for name in std::mem::take(&mut cmd.files) {
if is_link_not_fifo(host, &name) {
d.at(2, err, m::SYMLINK_IGNORED, &[&name]);
} else {
kept.push(name);
}
}
cmd.files = kept;
if cmd.files.is_empty() && total > 0 {
return clean(1);
}
}
for list in std::mem::take(&mut cmd.filelists) {
match util::names_from_file(host, &list, util_level, err) {
Some(names) => cmd.files.extend(names),
None => {
d.at(1, err, m::ERROR_READING, &[&list]);
return clean(1);
}
}
}
cmd.nb_input_files = cmd.files.len();
if cmd.recursive {
let names = std::mem::take(&mut cmd.files);
cmd.files = util::expand_names(host, names, cmd.follow_links, util_level, err);
}
match cmd.op {
Operation::List => return clean(list::list_files(cmd, host, out, err)),
Operation::Train => return finish(dibio::train_files(cmd, host, err), pause),
Operation::Test => {
cmd.prefs.test_mode = true;
cmd.out_file = Some(cli::NULMARK.to_vec());
cmd.prefs.remove_src = false;
}
Operation::Compress | Operation::Decompress => {}
}
if cmd.files.is_empty() {
if cmd.nb_input_files > 0 {
d.at(1, err, m::NO_INPUT_LEFT, &[]);
return clean(0);
}
cmd.files.push(STDINMARK.to_vec());
}
if cmd.files.len() == 1 && cmd.files[0] == STDINMARK && cmd.out_file.is_none() {
cmd.out_file = Some(STDOUTMARK.to_vec());
}
let has_stdin = cmd.files.iter().any(|f| f == STDINMARK);
if !cmd.force_stdin && has_stdin && (cmd.prefs.fake_stdin_console || host.stdin_is_tty()) {
d.at(1, err, m::STDIN_CONSOLE, &[]);
return clean(1);
}
if cmd.out_file.as_deref().is_none_or(|o| o == STDOUTMARK)
&& (cmd.prefs.fake_stdout_console || host.stdout_is_tty())
&& has_stdin
&& !cmd.force_stdout
&& cmd.op != Operation::Decompress
{
d.at(1, err, m::STDOUT_CONSOLE, &[]);
return clean(1);
}
let max_level = if cmd.ultra {
zstd_opts::MAX_CLEVEL
} else {
cli::CLEVEL_MAX
};
if cmd.level > max_level {
d.at(
2,
err,
m::LEVEL_REDUCED,
&[max_level.to_string().as_bytes()],
);
cmd.level = max_level;
}
if cmd.show_default_cparams && cmd.op == Operation::Decompress {
d.at(1, err, m::CPARAMS_DECOMPRESS, &[]);
return clean(1);
}
if cmd.dict_file.is_some() && cmd.patch_from.is_some() {
d.at(1, err, m::DICT_AND_PATCH, &[]);
return clean(1);
}
if cmd.patch_from.is_some() && cmd.files.len() > 1 {
d.at(1, err, m::PATCH_MULTIPLE, &[]);
return clean(1);
}
let has_stdout = cmd.out_file.as_deref() == Some(STDOUTMARK);
if has_stdout && d.level == 2 {
d.level = 1;
}
cmd.display_level = d.level;
let stderr_console = cmd.prefs.fake_stderr_console || host.stderr_is_tty();
if !stderr_console && cmd.prefs.progress != Progress::Always {
cmd.prefs.progress = Progress::Never;
}
if has_stdout && cmd.prefs.remove_src {
d.at(3, err, m::SRC_KEPT_STDOUT, &[]);
cmd.prefs.remove_src = false;
}
cmd.prefs.has_stdout_output = has_stdout;
cmd.prefs.nb_files_total = cmd.files.len();
cmd.prefs.has_stdin_input = has_stdin;
cmd.prefs.patch_from_mode = cmd.patch_from.is_some();
if cmd.prefs.memory_limit == 0 {
let wlog = cmd.cparam(CParamKey::WindowLog);
cmd.prefs.memory_limit = if wlog == 0 {
1 << cli::DEFAULT_MAX_WINDOW_LOG
} else {
1u32 << (wlog & 31)
};
}
if cmd.patch_from.is_some() {
cmd.dict_file = cmd.patch_from.clone();
}
if cmd.op != Operation::Compress {
if let Some(l) = cmd.limits {
cmd.prefs.memory_limit = cmd.prefs.memory_limit.min(1 << l.window_log_max);
}
return finish(fileio::decompress_files(cmd, host, out, err), pause);
}
if let Some(l) = cmd.limits
&& cmd.prefs.workers > l.workers
{
cmd.prefs.workers = l.workers;
}
let workers = cmd.prefs.workers;
if cmd.prefs.job_size as i32 != 0 && workers == 0 {
d.at(2, err, m::BLOCK_SIZE_USELESS, &[]);
}
if let Some(ovlog) = cmd.prefs.overlap_log
&& ovlog != 0
&& workers == 0
{
d.at(2, err, m::OVERLAP_LOG_USELESS, &[]);
}
if cmd.prefs.adapt && workers == 0 {
return exit(exm_throw(&d, err, 1, m::ADAPT_SINGLE_THREAD.as_bytes()));
}
if cmd.prefs.rsyncable && workers == 0 {
return exit(exm_throw(&d, err, 1, m::RSYNCABLE_SINGLE_THREAD.as_bytes()));
}
cmd.prefs.sparse = 0;
if cmd.prefs.adapt_min > cmd.level {
cmd.level = cmd.prefs.adapt_min;
}
if cmd.prefs.adapt_max < cmd.level {
cmd.level = cmd.prefs.adapt_max;
}
if cmd.show_default_cparams || d.level >= 4 {
let dict_size = cmd
.dict_file
.as_deref()
.map_or(0, |f| file_size(host, f).unwrap_or(u64::MAX));
for name in &cmd.files {
let size = file_size(host, name);
if cmd.show_default_cparams {
cli::print_default_cparams(err, name, size, dict_size, cmd.level);
}
if d.level >= 4 {
cli::print_actual_cparams(err, size, dict_size, cmd.level, cmd);
}
}
}
if d.level >= 4 {
display_compression_parameters(cmd, err);
}
finish(fileio::compress_files(cmd, host, out, err), pause)
}
fn display_compression_parameters(cmd: &Command, err: &mut dyn Write) {
let p = &cmd.prefs;
let n = |v: i64| v.to_string();
printf(
err,
m::DCP_FORMAT,
&[m::FORMAT_OPTIONS[p.ctype.index()].as_bytes()],
);
printf(err, m::DCP_SPARSE[p.sparse.clamp(0, 2) as usize], &[]);
if !p.dict_id {
printf(err, m::DCP_NO_DICT_ID, &[]);
}
printf(err, m::DCP_CHECKSUM[p.checksum.clamp(0, 2) as usize], &[]);
printf(
err,
m::DCP_BLOCK_SIZE,
&[n(i64::from(p.job_size as i32)).as_bytes()],
);
if p.adapt {
printf(
err,
m::DCP_ADAPT,
&[
n(p.adapt_min.into()).as_bytes(),
n(p.adapt_max.into()).as_bytes(),
],
);
}
printf(
err,
m::DCP_ROW_MATCH_FINDER[p.row_match_finder.value() as usize],
&[],
);
if p.rsyncable {
printf(err, m::DCP_RSYNCABLE, &[]);
}
if p.stream_size != 0 {
printf(
err,
m::DCP_STREAM_SIZE,
&[n((p.stream_size as u32).into()).as_bytes()],
);
}
if p.size_hint != 0 {
printf(err, m::DCP_SIZE_HINT, &[n(p.size_hint.into()).as_bytes()]);
}
if p.target_cblock_size != 0 {
printf(
err,
m::DCP_TARGET_CBLOCK,
&[n((p.target_cblock_size as u32).into()).as_bytes()],
);
}
printf(
err,
m::DCP_LITERALS[p.literal_compression.value() as usize],
&[],
);
let memory = if p.memory_limit != 0 {
p.memory_limit
} else {
128 << 20
};
printf(err, m::DCP_MEMORY, &[n(memory.into()).as_bytes()]);
printf(err, m::DCP_THREADS, &[n(p.workers.into()).as_bytes()]);
if p.exclude_compressed {
printf(err, m::DCP_EXCLUDE, &[]);
}
printf(
err,
m::DCP_CONTENT_SIZE,
&[if p.content_size { "" } else { m::DCP_NO }.as_bytes()],
);
}
#[cfg(test)]
mod tests;