use std::io::Read;
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
use dbmd_core::extract::{self, ExtractError};
use crate::cli::ExtractArgs;
use crate::context::Context;
use crate::error::{CliError, CliResult, ExitCode};
pub fn run(ctx: &Context, args: &ExtractArgs) -> CliResult {
let path = Path::new(&args.file);
let extracted = if std::env::var_os("DBMD_INTERNAL_EXTRACT_WORKER").is_some() {
#[cfg(debug_assertions)]
extraction_worker_test_hook();
extract::extract(path).map_err(map_extract_error)?
} else {
sandbox_extract(args)?
};
if ctx.json {
let json = serde_json::to_string_pretty(&extracted)
.map_err(|e| CliError::runtime(format!("failed to encode JSON: {e}")))?;
emit(&args.out, &json, true)
} else {
emit(&args.out, &extracted.text, false)
}
}
const EXTRACT_ADDRESS_SPACE_BYTES: u64 = 768 * 1024 * 1024;
const EXTRACT_CPU_SECONDS: u64 = 12;
const EXTRACT_ELAPSED_LIMIT: Duration = Duration::from_secs(20);
const EXTRACT_WORKER_OUTPUT_BYTES: u64 = 96 * 1024 * 1024;
fn sandbox_extract(args: &ExtractArgs) -> Result<extract::Extracted, CliError> {
#[cfg(not(unix))]
{
let _ = args;
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_SANDBOX_UNAVAILABLE",
"document extraction is disabled because this platform build cannot enforce memory and CPU limits",
));
}
#[cfg(unix)]
{
use std::os::unix::process::CommandExt as _;
let executable = std::env::current_exe().map_err(CliError::from)?;
#[cfg(debug_assertions)]
let cpu_seconds = std::env::var("DBMD_TEST_EXTRACT_CPU_SECONDS")
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|seconds| (1..=EXTRACT_CPU_SECONDS).contains(seconds))
.unwrap_or(EXTRACT_CPU_SECONDS);
#[cfg(not(debug_assertions))]
let cpu_seconds = EXTRACT_CPU_SECONDS;
#[cfg(debug_assertions)]
let skip_cpu_limit = std::env::var_os("DBMD_TEST_EXTRACT_SKIP_CPU_LIMIT").is_some();
#[cfg(not(debug_assertions))]
let skip_cpu_limit = false;
#[cfg(debug_assertions)]
let memory_limit_bytes = std::env::var("DBMD_TEST_EXTRACT_MEMORY_BYTES")
.ok()
.and_then(|value| value.parse::<u64>().ok())
.filter(|bytes| *bytes >= 16 * 1024 * 1024)
.unwrap_or(EXTRACT_ADDRESS_SPACE_BYTES);
#[cfg(not(debug_assertions))]
let memory_limit_bytes = EXTRACT_ADDRESS_SPACE_BYTES;
let mut command = Command::new(executable);
command
.arg("--json")
.arg("extract")
.arg("--")
.arg(&args.file)
.env("DBMD_INTERNAL_EXTRACT_WORKER", "1")
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
unsafe {
command.pre_exec(move || {
#[cfg(not(target_os = "macos"))]
let (memory_resource, memory) = (
libc::RLIMIT_AS,
libc::rlimit {
rlim_cur: memory_limit_bytes as libc::rlim_t,
rlim_max: memory_limit_bytes as libc::rlim_t,
},
);
#[cfg(not(target_os = "macos"))]
{
if libc::setrlimit(memory_resource, &memory) != 0 {
return Err(std::io::Error::last_os_error());
}
}
let cpu = libc::rlimit {
rlim_cur: cpu_seconds as libc::rlim_t,
rlim_max: cpu_seconds.saturating_add(1) as libc::rlim_t,
};
if !skip_cpu_limit && libc::setrlimit(libc::RLIMIT_CPU, &cpu) != 0 {
return Err(std::io::Error::last_os_error());
}
Ok(())
});
}
let mut child = command.spawn().map_err(|error| {
CliError::new(
ExitCode::Runtime,
"EXTRACT_SANDBOX_FAILED",
format!("could not start the extraction sandbox: {error}"),
)
})?;
let stdout = child.stdout.take().expect("piped extraction worker stdout");
let stderr = child.stderr.take().expect("piped extraction worker stderr");
let stdout_reader =
std::thread::spawn(move || read_worker_pipe(stdout, EXTRACT_WORKER_OUTPUT_BYTES));
let stderr_reader = std::thread::spawn(move || read_worker_pipe(stderr, 1024 * 1024));
let started = Instant::now();
let status = loop {
match child.try_wait() {
Ok(Some(status)) => break status,
Ok(None) if started.elapsed() < EXTRACT_ELAPSED_LIMIT => {
#[cfg(target_os = "macos")]
match macos_worker_resident_bytes(child.id()) {
Ok(bytes) if bytes > memory_limit_bytes => {
let _ = child.kill();
let _ = child.wait();
let _ = stdout_reader.join();
let _ = stderr_reader.join();
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_RESOURCE_LIMIT",
format!(
"document extraction exceeded the {} MiB resident-memory limit",
memory_limit_bytes / (1024 * 1024)
),
));
}
Ok(_) => {}
Err(error) => {
if let Ok(Some(status)) = child.try_wait() {
break status;
}
if error.kind() == std::io::ErrorKind::NotFound
|| error.raw_os_error() == Some(libc::ESRCH)
{
std::thread::sleep(Duration::from_millis(1));
continue;
}
let _ = child.kill();
let _ = child.wait();
let _ = stdout_reader.join();
let _ = stderr_reader.join();
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_SANDBOX_FAILED",
format!("could not inspect extraction worker memory: {error}"),
));
}
}
std::thread::sleep(Duration::from_millis(10));
}
Ok(None) => {
let _ = child.kill();
let _ = child.wait();
let _ = stdout_reader.join();
let _ = stderr_reader.join();
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_TIMEOUT",
format!(
"document extraction exceeded the {} second elapsed-time limit",
EXTRACT_ELAPSED_LIMIT.as_secs()
),
));
}
Err(error) => {
let _ = child.kill();
let _ = child.wait();
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_SANDBOX_FAILED",
format!("could not monitor the extraction sandbox: {error}"),
));
}
}
};
let stdout = stdout_reader
.join()
.map_err(|_| CliError::runtime("extraction stdout reader panicked"))??;
let stderr = stderr_reader
.join()
.map_err(|_| CliError::runtime("extraction stderr reader panicked"))??;
if !status.success() {
if let Ok(value) = serde_json::from_slice::<serde_json::Value>(&stderr) {
if let Some(error) = value.get("error") {
let code = error
.get("code")
.and_then(|value| value.as_str())
.unwrap_or("EXTRACT_PARSE_ERROR");
let message = error
.get("message")
.and_then(|value| value.as_str())
.unwrap_or("document extraction failed");
return Err(CliError::new(
ExitCode::Runtime,
stable_extract_code(code),
message,
));
}
}
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_RESOURCE_LIMIT",
format!(
"the extraction worker was terminated by a CPU, memory, or process failure ({status})"
),
));
}
serde_json::from_slice(&stdout).map_err(|error| {
CliError::new(
ExitCode::Runtime,
"EXTRACT_WORKER_PROTOCOL",
format!("invalid extraction worker response: {error}"),
)
})
}
}
#[cfg(target_os = "macos")]
fn macos_worker_resident_bytes(pid: u32) -> std::io::Result<u64> {
let mut info: libc::proc_taskinfo = unsafe { std::mem::zeroed() };
let expected = std::mem::size_of::<libc::proc_taskinfo>() as libc::c_int;
let written = unsafe {
libc::proc_pidinfo(
pid as libc::c_int,
libc::PROC_PIDTASKINFO,
0,
(&mut info as *mut libc::proc_taskinfo).cast(),
expected,
)
};
if written != expected {
let error = std::io::Error::last_os_error();
return Err(if error.raw_os_error().unwrap_or(0) == 0 {
std::io::Error::new(
std::io::ErrorKind::NotFound,
"proc_pidinfo returned no task information",
)
} else {
error
});
}
Ok(info.pti_resident_size)
}
#[cfg(debug_assertions)]
fn extraction_worker_test_hook() {
if let Some(bytes) = std::env::var("DBMD_TEST_EXTRACT_WORKER_ALLOCATE")
.ok()
.and_then(|value| value.parse::<usize>().ok())
{
let mut allocation = vec![0u8; bytes];
for page in allocation.chunks_mut(4096) {
page[0] = 0xA5;
}
std::hint::black_box(&mut allocation);
loop {
std::hint::spin_loop();
}
}
if std::env::var_os("DBMD_TEST_EXTRACT_WORKER_SPIN").is_some() {
loop {
std::hint::spin_loop();
}
}
}
fn read_worker_pipe<R: Read>(reader: R, max_bytes: u64) -> Result<Vec<u8>, CliError> {
let mut bytes = Vec::new();
reader
.take(max_bytes.saturating_add(1))
.read_to_end(&mut bytes)
.map_err(CliError::from)?;
if bytes.len() as u64 > max_bytes {
return Err(CliError::new(
ExitCode::Runtime,
"EXTRACT_WORKER_PROTOCOL",
"extraction worker output exceeded its transport cap",
));
}
Ok(bytes)
}
fn stable_extract_code(code: &str) -> &'static str {
match code {
"UNSUPPORTED_FORMAT" => "UNSUPPORTED_FORMAT",
"DOCUMENT_ENCRYPTED" => "DOCUMENT_ENCRYPTED",
"EXTRACT_PARSE_ERROR" => "EXTRACT_PARSE_ERROR",
"IO_ERROR" => "IO_ERROR",
_ => "EXTRACT_PARSE_ERROR",
}
}
fn emit(out: &Option<String>, content: &str, add_trailing_newline: bool) -> CliResult {
match out {
Some(path) => {
refuse_symlink_dest(path)?;
let mut body = content.to_string();
if add_trailing_newline && !body.ends_with('\n') {
body.push('\n');
}
dbmd_core::fsx::write_atomic(Path::new(path), body.as_bytes()).map_err(|e| {
CliError::new(
ExitCode::Runtime,
"IO_ERROR",
format!("failed to write {path}: {e}"),
)
})?;
Ok(())
}
None => {
use std::io::Write;
let stdout = std::io::stdout();
let mut lock = stdout.lock();
let res = if add_trailing_newline {
writeln!(lock, "{content}")
} else {
write!(lock, "{content}")
};
match res {
Ok(()) => Ok(()),
Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => Ok(()),
Err(e) => Err(CliError::new(
ExitCode::Runtime,
"IO_ERROR",
format!("write failed: {e}"),
)),
}
}
}
}
fn refuse_symlink_dest(path: &str) -> Result<(), CliError> {
use std::path::Component;
let refuse = |p: &Path| {
CliError::new(
ExitCode::Runtime,
"OUT_IS_SYMLINK",
format!(
"refusing to write {path}: the path is reached through a symlink ({})",
p.display()
),
)
.with_hint(
"extract --out will not follow a symlink (it could overwrite a file elsewhere); \
remove the symlink or choose a destination with no symlinked component",
)
};
let inspect_io_err = |p: &Path, e: std::io::Error| {
CliError::new(
ExitCode::Runtime,
"IO_ERROR",
format!("failed to inspect {}: {e}", p.display()),
)
};
let leaf = Path::new(path);
match std::fs::symlink_metadata(leaf) {
Ok(meta) if meta.file_type().is_symlink() => return Err(refuse(leaf)),
Ok(_) => {}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => return Err(inspect_io_err(leaf, e)),
}
let parent = match leaf.parent() {
Some(p) if !p.as_os_str().is_empty() => p,
_ => return Ok(()),
};
let mut existing = parent.to_path_buf();
let exists = |p: &Path| match std::fs::symlink_metadata(p) {
Ok(_) => Ok(true),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
Err(e) => Err(e),
};
loop {
match exists(&existing) {
Ok(true) => break,
Ok(false) => {
if !existing.pop() || existing.as_os_str().is_empty() {
return Ok(());
}
}
Err(e) => return Err(inspect_io_err(&existing, e)),
}
}
let mut real = std::path::PathBuf::new();
let mut lexical = std::path::PathBuf::new();
let mut on_real_ground = false;
if existing.is_relative() {
real = match std::env::current_dir() {
Ok(cwd) => match cwd.canonicalize() {
Ok(c) => c,
Err(e) => return Err(inspect_io_err(&cwd, e)),
},
Err(e) => return Err(inspect_io_err(Path::new("."), e)),
};
on_real_ground = true;
}
for comp in existing.components() {
match comp {
Component::Prefix(_) | Component::RootDir => {
real.push(comp.as_os_str());
lexical.push(comp.as_os_str());
}
Component::CurDir => {}
Component::ParentDir => {
real.pop();
lexical.pop();
}
Component::Normal(name) => {
lexical.push(name);
let probe = real.join(name);
match std::fs::symlink_metadata(&probe) {
Ok(meta) if meta.file_type().is_symlink() => {
if on_real_ground {
return Err(refuse(&lexical));
}
match probe.canonicalize() {
Ok(c) => real = c,
Err(e) => return Err(inspect_io_err(&probe, e)),
}
}
Ok(meta) => {
if meta.is_dir() {
on_real_ground = true;
}
real = probe;
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => return Err(inspect_io_err(&probe, e)),
}
}
}
}
Ok(())
}
fn map_extract_error(err: ExtractError) -> CliError {
match &err {
ExtractError::UnsupportedFormat(_) => CliError::new(
ExitCode::Runtime,
err.code(),
err.to_string(),
)
.with_hint(
"supported document types: .pdf, .docx, .xlsx/.xlsm/.xlsb/.ods, .epub, .html/.htm/.xhtml (detected by extension)",
),
ExtractError::Encrypted(_) => CliError::new(ExitCode::Runtime, err.code(), err.to_string())
.with_hint("the document is password-protected; dbmd extract cannot open it"),
ExtractError::Parse { .. } => CliError::new(ExitCode::Runtime, err.code(), err.to_string()),
ExtractError::Io(_) => CliError::new(ExitCode::Runtime, "IO_ERROR", err.to_string()),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[cfg(unix)]
fn out_symlink_is_refused_and_target_untouched() {
let tmp = tempfile::tempdir().unwrap();
let victim = tmp.path().join("victim.conf");
std::fs::write(&victim, "SENSITIVE-ORIGINAL\n").unwrap();
let link = tmp.path().join("innocent-output.txt");
std::os::unix::fs::symlink(&victim, &link).unwrap();
let out = Some(link.to_string_lossy().into_owned());
let err = emit(&out, "POISONED-BYTES-FROM-DOCUMENT", false)
.expect_err("a symlinked --out must be refused");
assert_eq!(err.code, "OUT_IS_SYMLINK", "got {err:?}");
assert_eq!(
std::fs::read_to_string(&victim).unwrap(),
"SENSITIVE-ORIGINAL\n",
"the symlink target must not be overwritten",
);
assert!(
std::fs::symlink_metadata(&link)
.unwrap()
.file_type()
.is_symlink(),
"the --out path must remain a symlink",
);
}
#[test]
#[cfg(unix)]
fn out_through_symlinked_parent_is_refused_and_target_untouched() {
let tmp = tempfile::tempdir().unwrap();
let store = tmp.path().join("store");
std::fs::create_dir(&store).unwrap();
let external = tmp.path().join("external");
std::fs::create_dir(&external).unwrap();
let victim = external.join("victim.txt");
std::fs::write(&victim, "ORIGINAL_SECRET\n").unwrap();
let linkdir = store.join("linkdir");
std::os::unix::fs::symlink(&external, &linkdir).unwrap();
let out_path = linkdir.join("victim.txt");
let out = Some(out_path.to_string_lossy().into_owned());
let err = emit(&out, "POISONED_BY_EXTRACT", false)
.expect_err("a --out reached through a symlinked parent must be refused");
assert_eq!(err.code, "OUT_IS_SYMLINK", "got {err:?}");
assert_eq!(
std::fs::read_to_string(&victim).unwrap(),
"ORIGINAL_SECRET\n",
"the symlinked-parent target must not be overwritten",
);
}
#[test]
#[cfg(unix)]
fn out_through_symlinked_parent_with_real_subdir_is_refused() {
let tmp = tempfile::tempdir().unwrap();
let store = tmp.path().join("store");
std::fs::create_dir(&store).unwrap();
let external = tmp.path().join("external");
let external_sub = external.join("sub");
std::fs::create_dir_all(&external_sub).unwrap();
let victim = external_sub.join("victim.txt");
std::fs::write(&victim, "ORIGINAL_SECRET\n").unwrap();
let linkdir = store.join("linkdir");
std::os::unix::fs::symlink(&external, &linkdir).unwrap();
let out_path = linkdir.join("sub").join("victim.txt");
let out = Some(out_path.to_string_lossy().into_owned());
let err = emit(&out, "POISONED_BY_EXTRACT", false).expect_err(
"a --out reached through a symlinked parent (with a real subdir below \
the link) must be refused",
);
assert_eq!(err.code, "OUT_IS_SYMLINK", "got {err:?}");
assert_eq!(
std::fs::read_to_string(&victim).unwrap(),
"ORIGINAL_SECRET\n",
"the deep symlinked-parent target must not be overwritten",
);
}
#[test]
fn out_into_real_nested_subdir_is_written() {
let tmp = tempfile::tempdir().unwrap();
let nested = tmp.path().join("a").join("b").join("c");
std::fs::create_dir_all(&nested).unwrap();
let dest = nested.join("out.txt");
let out = Some(dest.to_string_lossy().into_owned());
emit(&out, "deep but real", false).expect("a real deep-nested --out must succeed");
assert_eq!(std::fs::read_to_string(&dest).unwrap(), "deep but real");
}
#[test]
fn out_regular_file_is_written() {
let tmp = tempfile::tempdir().unwrap();
let dest = tmp.path().join("out.txt");
let out = Some(dest.to_string_lossy().into_owned());
emit(&out, "hello extracted text", false).expect("a regular --out must succeed");
assert_eq!(
std::fs::read_to_string(&dest).unwrap(),
"hello extracted text",
);
emit(&out, "second write", false).expect("overwriting a regular file is allowed");
assert_eq!(std::fs::read_to_string(&dest).unwrap(), "second write");
}
}