use std::ffi::OsString;
use std::process::ExitCode;
use std::sync::Arc;
use qld::args::{OutputCompleteHook, ParseOutcome};
use qld::diag::{Diagnostic, DiagnosticSink, Stderr};
fn main() -> ExitCode {
tune_allocator();
let args: Vec<OsString> = std::env::args_os().collect();
#[cfg(unix)]
if let Some(child) = fork::Child::from_env() {
let code = run(&args, Launch::Child(child));
child.notify(code);
return ExitCode::from(code);
}
ExitCode::from(run(&args, Launch::MayFork))
}
#[cfg(all(unix, target_env = "gnu"))]
fn tune_allocator() {
const M_TRIM_THRESHOLD: i32 = -1;
const M_TOP_PAD: i32 = -2;
unsafe extern "C" {
fn mallopt(param: i32, value: i32) -> i32;
}
unsafe {
mallopt(M_TOP_PAD, 64 * 1024 * 1024);
mallopt(M_TRIM_THRESHOLD, 128 * 1024 * 1024);
}
}
#[cfg(not(all(unix, target_env = "gnu")))]
fn tune_allocator() {}
enum Launch {
MayFork,
#[cfg(unix)]
Child(fork::Child),
}
fn run(args: &[OsString], launch: Launch) -> u8 {
let diagnostics = Arc::new(Stderr::new(qld::PROGRAM_NAME));
let hook = match launch {
Launch::MayFork => None,
#[cfg(unix)]
Launch::Child(child) => {
let sink = Arc::clone(&diagnostics);
Some(OutputCompleteHook::new(move || {
sink.flush();
child.notify(u8::from(sink.error_count() != 0));
}))
}
};
let result = run_link(args, hook, &diagnostics);
diagnostics.flush();
match result {
Ok(code) if diagnostics.error_count() == 0 => code,
Ok(_) => 1,
Err(qld::Error::Reported { .. }) => 1,
Err(error) => {
eprintln!("{}: error: {error}", qld::PROGRAM_NAME);
1
}
}
}
fn run_link(
args: &[OsString],
hook: Option<OutputCompleteHook>,
diagnostics: &Stderr,
) -> qld::Result<u8> {
let mut options = match qld::parse_gnu(args)? {
ParseOutcome::Help => {
print!("{}", qld::args::usage());
return Ok(0);
}
ParseOutcome::Version => {
println!("{}", qld::version_line());
return Ok(0);
}
ParseOutcome::Link(options) => options,
};
diagnostics.configure(&options);
match hook {
None => {
#[cfg(unix)]
if options.fork
&& fork::allowed(args, &options)
&& let Some(code) = fork::link_in_child(args)
{
return Ok(code);
}
}
Some(hook) => options.on_output_complete = Some(hook),
}
options.exit_on_plugin_fatal = true;
for warning in &options.warnings {
diagnostics.emit(Diagnostic::warning(warning.clone()));
}
qld::link(&options, diagnostics)?;
Ok(0)
}
#[cfg(unix)]
mod fork {
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Read, Write};
use std::net::Shutdown;
use std::os::fd::{AsFd, BorrowedFd, OwnedFd};
use std::os::unix::fs::{FileTypeExt, MetadataExt};
use std::os::unix::net::UnixStream;
use std::os::unix::process::{CommandExt, ExitStatusExt};
use std::path::PathBuf;
use std::process::{Command, ExitStatus, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::thread::JoinHandle;
use qld::LinkOptions;
const ENV: &str = "QLD_FORK_CHILD";
pub fn allowed(args: &[OsString], options: &LinkOptions) -> bool {
let paths = [Some(options.output_path()), options.map_file.clone()];
!args
.iter()
.any(|arg| names_system_path(arg.as_encoded_bytes()))
&& !paths
.iter()
.flatten()
.any(|path| path.starts_with("/dev") || path.starts_with("/proc"))
}
fn names_system_path(arg: &[u8]) -> bool {
[&b"/dev/"[..], b"/proc/"].iter().any(|dir| {
(0..arg.len()).any(|at| {
let (before, rest) = arg.split_at(at);
rest.starts_with(dir)
&& (before.is_empty()
|| before.ends_with(b"=")
|| before.ends_with(b"@")
|| (before.starts_with(b"-") && !before.contains(&b'/')))
})
})
}
struct Stream {
relay: bool,
id: (u64, u64),
}
impl Stream {
fn of(fd: BorrowedFd<'_>) -> Option<Self> {
let metadata = File::from(fd.try_clone_to_owned().ok()?).metadata().ok()?;
let kind = metadata.file_type();
Some(Self {
relay: kind.is_fifo() || kind.is_socket(),
id: (metadata.dev(), metadata.ino()),
})
}
}
pub fn link_in_child(args: &[OsString]) -> Option<u8> {
let stdout = Stream::of(io::stdout().as_fd())?;
let stderr = Stream::of(io::stderr().as_fd())?;
let shared = stdout.relay && stderr.relay && stdout.id == stderr.id;
let (status, status_child) = UnixStream::pair().ok()?;
let mut command = Command::new(program()?);
if let Some(arg0) = args.first() {
command.arg0(arg0);
}
command
.args(args.iter().skip(1))
.env(
ENV,
format!(
"{}:{}{}",
std::process::id(),
if stdout.relay { "o" } else { "" },
if stderr.relay { "e" } else { "" }
),
)
.stdin(Stdio::from(OwnedFd::from(status_child)));
let mut relays = Vec::new();
if stdout.relay {
let child = relay(io::stdout().as_fd(), &mut relays)?;
if shared {
command.stderr(Stdio::from(OwnedFd::from(child.try_clone().ok()?)));
}
command.stdout(Stdio::from(OwnedFd::from(child)));
}
if stderr.relay && !shared {
let child = relay(io::stderr().as_fd(), &mut relays)?;
command.stderr(Stdio::from(OwnedFd::from(child)));
}
let spawned = command.spawn();
drop(command);
let Ok(mut child) = spawned else {
join(relays);
return None;
};
let _ = status.shutdown(Shutdown::Write);
let reported = read_status(&status);
join(relays);
Some(match reported {
Some(code) => code,
None => exit_code(child.wait()),
})
}
fn program() -> Option<PathBuf> {
let path = std::env::current_exe().ok()?;
#[cfg(any(target_os = "linux", target_os = "android"))]
{
let id = |path: &std::path::Path| {
std::fs::metadata(path)
.ok()
.map(|metadata| (metadata.dev(), metadata.ino()))
};
let proc = PathBuf::from("/proc/self/exe");
if let Some(this) = id(&proc)
&& id(&path) != Some(this)
{
return Some(proc);
}
}
Some(path)
}
fn relay(target: BorrowedFd<'_>, relays: &mut Vec<JoinHandle<()>>) -> Option<UnixStream> {
let mut target = File::from(target.try_clone_to_owned().ok()?);
let (mut source, child) = UnixStream::pair().ok()?;
let thread = std::thread::Builder::new()
.name("qld-relay".into())
.spawn(move || {
let mut buffer = vec![0u8; 64 << 10];
loop {
match source.read(&mut buffer) {
Ok(0) => return,
Ok(n) => {
if target.write_all(buffer.get(..n).unwrap_or(&[])).is_err() {
return;
}
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(_) => return,
}
}
})
.ok()?;
relays.push(thread);
Some(child)
}
fn join(relays: Vec<JoinHandle<()>>) {
for relay in relays {
let _ = relay.join();
}
}
fn read_status(mut status: &UnixStream) -> Option<u8> {
let mut byte = [0u8; 1];
loop {
match status.read(&mut byte) {
Ok(1) => return Some(byte[0]),
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Ok(_) | Err(_) => return None,
}
}
}
fn exit_code(status: io::Result<ExitStatus>) -> u8 {
let program = qld::PROGRAM_NAME;
match status {
Ok(status) => {
if let Some(code) = status.code() {
return u8::try_from(code & 0xff).unwrap_or(1);
}
let Some(signal) = status.signal() else {
return 1;
};
let core = if status.core_dumped() {
" (core dumped)"
} else {
""
};
eprintln!("{program}: error: the link process was killed by signal {signal}{core}");
u8::try_from(signal.saturating_add(128)).unwrap_or(1)
}
Err(error) => {
eprintln!("{program}: error: cannot wait for the link process: {error}");
1
}
}
}
#[derive(Clone, Copy)]
pub struct Child {
relay_stdout: bool,
relay_stderr: bool,
}
impl Child {
pub fn from_env() -> Option<Self> {
let value = std::env::var(ENV).ok()?;
let (pid, flags) = value.split_once(':')?;
if pid.parse::<u32>().ok()? != std::os::unix::process::parent_id() {
return None;
}
let stdin = File::from(io::stdin().as_fd().try_clone_to_owned().ok()?);
if !stdin.metadata().ok()?.file_type().is_socket() {
return None;
}
Some(Self {
relay_stdout: flags.contains('o'),
relay_stderr: flags.contains('e'),
})
}
pub fn notify(self, code: u8) {
static SENT: AtomicBool = AtomicBool::new(false);
if SENT.swap(true, Ordering::AcqRel) {
return;
}
let _ = io::stdout().flush();
let _ = io::stderr().flush();
if self.relay_stdout {
shut_down(io::stdout().as_fd());
}
if self.relay_stderr {
shut_down(io::stderr().as_fd());
}
if let Ok(fd) = io::stdin().as_fd().try_clone_to_owned() {
let _ = UnixStream::from(fd).write_all(&[code]);
}
}
}
fn shut_down(fd: BorrowedFd<'_>) {
if let Ok(fd) = fd.try_clone_to_owned() {
let _ = UnixStream::from(fd).shutdown(Shutdown::Write);
}
}
#[cfg(test)]
mod tests {
use super::names_system_path;
#[test]
fn system_paths_in_arguments() {
for arg in [
"/dev/stdout",
"/proc/self/fd/1",
"-o/dev/stdout",
"--output=/dev/stdout",
"-Map=/dev/stderr",
"@/dev/stdin",
"-T/dev/stdin",
] {
assert!(names_system_path(arg.as_bytes()), "{arg}");
}
for arg in [
"/home/me/dev/out",
"-L/home/me/dev/lib",
"--output=/home/me/proc/a.out",
"dev/stdout",
"-lproc",
"",
] {
assert!(!names_system_path(arg.as_bytes()), "{arg}");
}
}
}
}