diff --git a/src/main.rs b/src/main.rs index 73e4a3b..ace97bd 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,6 +11,9 @@ use std::os::fd::{BorrowedFd, FromRawFd}; use std::os::unix::io::AsRawFd; use std::ptr; +static mut USE_KERNEL_ORDERING: bool = false; +static mut PLAIN_MODE: bool = false; + #[derive(Debug, Default)] struct Ptys { out_slave: libc::c_int, @@ -39,20 +42,31 @@ fn handle_data(mut stream: StdioStream) -> Result<(), bool> { Ok(0) => Err(false), // EOF Ok(n) => { let bufc = &buf[..n]; - for c in bufc { - stream.buffer.push(*c); - if *c == 0x0a as u8 { - let format = String::from("%H:%M:%S.%6f") - + match (stream.is_tty, stream.var) { - (true, StdioStreamType::StdOut) => " \x1b[32m(O)\x1b[0m ", - (false, StdioStreamType::StdOut) => " (O) ", - (true, StdioStreamType::StdErr) => " \x1b[31m(E)\x1b[0m ", - (false, StdioStreamType::StdErr) => " (E) ", - }; - let now = Local::now().format(&format).to_string(); - std::io::stdout().write_all(now.as_bytes()).unwrap(); - std::io::stdout().write_all(&stream.buffer).unwrap(); - stream.buffer.clear(); + if unsafe {PLAIN_MODE} { + std::io::stdout().write_all(bufc).unwrap(); + } else { + for c in bufc { + stream.buffer.push(*c); + if *c == 0x0a as u8 { + let format = String::from("%H:%M:%S.%6f") + + if unsafe{USE_KERNEL_ORDERING} { + match stream.is_tty { + true => " \x1b[34m(?)\x1b[0m ", + false => " (?) ", + } + } else { + match (stream.is_tty, stream.var) { + (true, StdioStreamType::StdOut) => " \x1b[32m(O)\x1b[0m ", + (false, StdioStreamType::StdOut) => " (O) ", + (true, StdioStreamType::StdErr) => " \x1b[31m(E)\x1b[0m ", + (false, StdioStreamType::StdErr) => " (E) ", + } + }; + let now = Local::now().format(&format).to_string(); + std::io::stdout().write_all(now.as_bytes()).unwrap(); + std::io::stdout().write_all(&stream.buffer).unwrap(); + stream.buffer.clear(); + } } } std::io::stdout().flush().unwrap(); @@ -89,13 +103,14 @@ fn open_ptys(ptys: &mut Ptys) { assert!(ptys.err_master != 0); } +static mut IS_TTY: bool = false; + fn handle_poll( poll_fd: &PollFd<'_>, var: StdioStreamType, master_file: &File, mut buf: &mut Vec, ) -> Result<(), bool> { - let is_tty = unsafe { libc::isatty(libc::STDOUT_FILENO) != 0 }; let revents = poll_fd.revents(); if let Some(revents) = revents { if revents.contains(PollFlags::POLLHUP) { @@ -105,25 +120,32 @@ fn handle_poll( return handle_data(StdioStream { var: var, file: &master_file, - is_tty: is_tty, + is_tty: unsafe { IS_TTY }, buffer: &mut buf, }); } + Err(false) + } else { + Err(false) } - Ok(()) } -fn spawn_child(ptys: &Ptys) { +fn spawn_child(ptys: &Ptys, user_cmd: String) { let cmd = CString::new("/bin/bash").unwrap(); - let bcmd = std::env::args().skip(1).collect::>().join(" "); - println!("Cmd: '{}'", bcmd); - + let bcmd = user_cmd; + if !unsafe {PLAIN_MODE} { + println!("Cmd: '{}'", bcmd); + } setsid().expect("Failed to create new session"); unsafe { libc::close(ptys.out_master.as_raw_fd()) }; unsafe { libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDOUT_FILENO); - libc::dup2(ptys.err_slave.as_raw_fd(), libc::STDERR_FILENO); + if USE_KERNEL_ORDERING { + libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDERR_FILENO); + } else { + libc::dup2(ptys.err_slave.as_raw_fd(), libc::STDERR_FILENO); + } // There's no real reason to dup the stdin. // It's not being used anyway. // ---------------------------------------------------------- @@ -173,8 +195,11 @@ fn main_loop(child: Pid, ptys: &Ptys) { ), ]; - let mut out_buf: Vec = Vec::new(); - let mut err_buf: Vec = Vec::new(); + let mut out_buf: Vec = Vec::with_capacity(65536); + let mut err_buf: Vec = Vec::with_capacity(65536); + + let mut writing_out = false; + let mut writing_err = false; loop { let res = poll(&mut poll_fds, PollTimeout::ZERO); @@ -183,52 +208,143 @@ fn main_loop(child: Pid, ptys: &Ptys) { std::process::exit(-1); } - match handle_poll( - &poll_fds[1], - StdioStreamType::StdErr, - &master_file_err, - &mut err_buf, - ) { - Ok(_) => {} - Err(true) => break, - Err(false) => continue, - } - match handle_poll( - &poll_fds[0], - StdioStreamType::StdOut, - &master_file_out, - &mut out_buf, - ) { - Ok(_) => {} - Err(true) => break, - Err(false) => continue, + assert!(res != Ok(2), "Race condition between stdio and stderr"); + + if unsafe{USE_KERNEL_ORDERING} { + match handle_poll( + &poll_fds[0], + StdioStreamType::StdOut, + &master_file_out, + &mut out_buf, + ) { + Ok(_) => {} + Err(true) => break, + Err(false) => {} + } + } else { + if !writing_out { + match handle_poll( + &poll_fds[1], + StdioStreamType::StdErr, + &master_file_err, + &mut err_buf, + ) { + Ok(_) => { + writing_err = true; + writing_out = false; + } + Err(true) => break, + Err(false) => { + writing_err = false; + writing_out = true; + } + } + } + if !writing_err { + match handle_poll( + &poll_fds[0], + StdioStreamType::StdOut, + &master_file_out, + &mut out_buf, + ) { + Ok(_) => { + writing_err = false; + writing_out = true; + } + Err(true) => break, + Err(false) => { + writing_err = true; + writing_out = false; + } + } + } } - if res.ok().unwrap() == 0 && unsafe { kill(child.as_raw(), 0) } != 0 { - // The poll might've timed out and in the short time between - // the poll timeout and the check if the process is alive, - // there might be pending data (unless the kernel deletes it...) - // I have no evidence supporting this theory, though. - // ------------------------------------------------------------- - // let mut buf = [0u8; 64]; - // match master_file.read(&mut buf) { - // Ok(e) => println!("Ok: {}", e), - // Err(_e) => {/*println!("Err: {}", e)*/} - // } - break; + if res.ok().unwrap() == 0 { + writing_out = false; + writing_err = false; + if unsafe { kill(child.as_raw(), 0) } != 0 { + // The poll might've timed out and in the short time between + // the poll timeout and the check if the process is alive, + // there might be pending data (unless the kernel deletes it...) + // I have no evidence supporting this theory, though. + // ------------------------------------------------------------- + // let mut buf = [0u8; 64]; + // match master_file.read(&mut buf) { + // Ok(e) => println!("Ok: {}", e), + // Err(_e) => {/*println!("Err: {}", e)*/} + // } + break; + } } waitpid(child, Some(WaitPidFlag::WNOHANG)).ok(); } } +use std::ffi::OsString; + +#[derive(Debug)] +struct Args { + kernel_order: bool, + plain: bool, + command: String, +} + +fn parse_args() -> Args { + let osargs = std::env::args_os(); + let args = osargs.collect::>(); + let mut args = args.into_iter().skip(1); // skip program name + let mut kernel_order = false; + let mut plain = false; + let mut command_parts: Vec = Vec::new(); + + while let Some(arg) = args.next() { + if arg == "--kernel-order" { + kernel_order = true; + } else if arg == "--plain" { + plain = true; + } else if arg == "-c" || arg == "--command" { + command_parts.extend(args); + break; + } else { + command_parts.push(arg); + command_parts.extend(args); + break; + } + } + + let command = command_parts + .into_iter() + .map(|s| s.to_string_lossy().into_owned()) + .collect::>() + .join(" "); + + let parsed = Args { + kernel_order, + command, + plain, + }; + + parsed +} + + fn main() { + let args = parse_args(); + let cmd = args.command; + + unsafe { + IS_TTY = libc::isatty(libc::STDOUT_FILENO) != 0; + USE_KERNEL_ORDERING = args.kernel_order; + PLAIN_MODE = args.plain; + }; let mut ptys = Ptys::default(); open_ptys(&mut ptys); match unsafe { fork() } { Ok(ForkResult::Child) => { - spawn_child(&ptys); + spawn_child(&ptys, cmd); } Ok(ForkResult::Parent { child }) => {