diff --git a/src/main.rs b/src/main.rs index 2fbe819..73e4a3b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,7 @@ use chrono::Local; use libc::kill; use nix::poll::{PollFd, PollFlags, PollTimeout, poll}; use nix::sys::wait::{WaitPidFlag, waitpid}; -use nix::unistd::{ForkResult, execvp, fork, setsid}; +use nix::unistd::{ForkResult, Pid, execvp, fork, setsid}; use std::ffi::CString; use std::fs::File; use std::io::Read; @@ -32,11 +32,11 @@ struct StdioStream<'a> { buffer: &'a mut Vec, } -fn handle_data(mut stream: StdioStream) -> Result<(), ()> { +fn handle_data(mut stream: StdioStream) -> Result<(), bool> { let mut buf = [0u8; 65536]; let res = stream.file.read(&mut buf); match res { - Ok(0) => Err(()), // EOF + Ok(0) => Err(false), // EOF Ok(n) => { let bufc = &buf[..n]; for c in bufc { @@ -60,13 +60,12 @@ fn handle_data(mut stream: StdioStream) -> Result<(), ()> { } Err(e) => { eprintln!("Read error: {}", e); - Err(()) + Err(true) } } } -fn main() { - let mut ptys = Ptys::default(); +fn open_ptys(ptys: &mut Ptys) { unsafe { libc::openpty( &mut ptys.out_slave, @@ -83,137 +82,157 @@ fn main() { ptr::null_mut(), ); } + assert!(ptys.out_slave != 0); assert!(ptys.out_master != 0); assert!(ptys.err_slave != 0); assert!(ptys.err_master != 0); +} +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) { + return Err(true); + } + if revents.contains(PollFlags::POLLIN) { + return handle_data(StdioStream { + var: var, + file: &master_file, + is_tty: is_tty, + buffer: &mut buf, + }); + } + } + Ok(()) +} - match unsafe { fork() } { - Ok(ForkResult::Child) => { - let cmd = CString::new("/bin/bash").unwrap(); - let bcmd = std::env::args().skip(1).collect::>().join(" "); - 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); - // There's no real reason to dup the stdin. - // It's not being used anyway. - // ---------------------------------------------------------- - // libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDIN_FILENO); - } +fn spawn_child(ptys: &Ptys) { + let cmd = CString::new("/bin/bash").unwrap(); + let bcmd = std::env::args().skip(1).collect::>().join(" "); + 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); + // There's no real reason to dup the stdin. + // It's not being used anyway. + // ---------------------------------------------------------- + // libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDIN_FILENO); + } + + unsafe { libc::close(ptys.out_slave.as_raw_fd()) }; + unsafe { libc::close(ptys.err_slave.as_raw_fd()) }; + + let args = [ + cmd.clone(), + CString::new("-c").unwrap(), + CString::new(bcmd).unwrap(), + ]; + execvp(&cmd, &args).expect("execvp failed"); +} - unsafe { libc::close(ptys.out_slave.as_raw_fd()) }; - unsafe { libc::close(ptys.err_slave.as_raw_fd()) }; +fn main_loop(child: Pid, ptys: &Ptys) { + // TODO: The kernel seems to garbage collect the pty + // & the associated data when the only holder of the slave end dies, + // resulting in errors reading it. So we keep a copy around. + // I don't actually know if this is the correct way of doing things. + // But it fixed the issues, so... + // ----------------------------------------------------------------- + // unsafe { libc::close(slave.as_raw_fd()) }; + + let master_file_out = unsafe { File::from_raw_fd(ptys.out_master) }; + let master_file_err = unsafe { File::from_raw_fd(ptys.err_master) }; + + // In some racy cases, the read call might get stuck... + // Maybe it'd be better to switch to non-blocking reads instead of + // using poll. + // --------------------------------------------------------------- + // unsafe { + // let flags = fcntl(ptys.out_master, F_GETFL, 0); + // fcntl(ptys.out_master, F_SETFL, flags | O_NONBLOCK); + // } + + let mut poll_fds = [ + PollFd::new( + unsafe { BorrowedFd::borrow_raw(ptys.out_master) }, + PollFlags::POLLIN | PollFlags::POLLHUP, + ), + PollFd::new( + unsafe { BorrowedFd::borrow_raw(ptys.err_master) }, + PollFlags::POLLIN | PollFlags::POLLHUP, + ), + ]; + + let mut out_buf: Vec = Vec::new(); + let mut err_buf: Vec = Vec::new(); + + loop { + let res = poll(&mut poll_fds, PollTimeout::ZERO); + if res.is_err() { + eprintln!("Poll failed"); + std::process::exit(-1); + } - let args = [ - cmd.clone(), - CString::new("-c").unwrap(), - CString::new(bcmd).unwrap(), - ]; - execvp(&cmd, &args).expect("execvp failed"); + 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, } - Ok(ForkResult::Parent { child }) => { - let _ = child; - // TODO: The kernel seems to garbage collect the pty - // & the associated data when the only holder of the slave end dies, - // resulting in errors reading it. So we keep a copy around. - // I don't actually know if this is the correct way of doing things. - // But it fixed the issues, so... - // ----------------------------------------------------------------- - // unsafe { libc::close(slave.as_raw_fd()) }; - - let master_file_out = unsafe { File::from_raw_fd(ptys.out_master) }; - let master_file_err = unsafe { File::from_raw_fd(ptys.err_master) }; - - // In some racy cases, the read call might get stuck... - // Maybe it'd be better to switch to non-blocking reads instead of - // using poll. - // --------------------------------------------------------------- - // unsafe { - // let flags = fcntl(ptys.out_master, F_GETFL, 0); - // fcntl(ptys.out_master, F_SETFL, flags | O_NONBLOCK); + 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; + } - let mut poll_fds = [ - PollFd::new( - unsafe { BorrowedFd::borrow_raw(ptys.out_master) }, - PollFlags::POLLIN | PollFlags::POLLHUP, - ), - PollFd::new( - unsafe { BorrowedFd::borrow_raw(ptys.err_master) }, - PollFlags::POLLIN | PollFlags::POLLHUP, - ), - ]; - - let mut out_buf: Vec = Vec::new(); - let mut err_buf: Vec = Vec::new(); - - loop { - let res = poll(&mut poll_fds, PollTimeout::ZERO); - if res.is_err() { - eprintln!("Poll failed"); - std::process::exit(-1); - } - - let revents = poll_fds[0].revents(); - if let Some(revents) = revents { - if revents.contains(PollFlags::POLLHUP) { - break; - } - if revents.contains(PollFlags::POLLIN) { - match handle_data(StdioStream { - var: StdioStreamType::StdOut, - file: &master_file_out, - is_tty: is_tty, - buffer: &mut out_buf, - }) { - Ok(_) => {} - Err(_) => break, - } - } - } + waitpid(child, Some(WaitPidFlag::WNOHANG)).ok(); + } +} - let revents = poll_fds[1].revents(); - if let Some(revents) = revents { - if revents.contains(PollFlags::POLLHUP) { - break; - } - if revents.contains(PollFlags::POLLIN) { - match handle_data(StdioStream { - var: StdioStreamType::StdErr, - file: &master_file_err, - is_tty: is_tty, - buffer: &mut err_buf, - }) { - Ok(_) => {} - Err(_) => break, - } - } - } +fn main() { + let mut ptys = Ptys::default(); + open_ptys(&mut ptys); - 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; - } + match unsafe { fork() } { + Ok(ForkResult::Child) => { + spawn_child(&ptys); + } - waitpid(child, Some(WaitPidFlag::WNOHANG)).ok(); - } + Ok(ForkResult::Parent { child }) => { + main_loop(child, &ptys); } Err(e) => { @@ -221,5 +240,4 @@ fn main() { std::process::exit(-1); } } - println!(""); }