itycodes
/
stdgather
1
0
Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

Minor cleanup

Browse Source 
Deduplicated some code into functions. Behavior should be equivalent to
the behavior of the last commit.
master
itycodes 3 weeks ago
parent e90adb6c0d
commit 20ffb15281
1 changed files with 140 additions and 122 deletions
Show Stats Download Patch File Download Diff File

262
src/main.rs
Unescape View File

@ -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<u8>,
}
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<u8>,
) -> 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::<Vec<String>>().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::<Vec<String>>().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<u8> = Vec::new();
let mut err_buf: Vec<u8> = 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<u8> = Vec::new();
let mut err_buf: Vec<u8> = 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!("");
}

Write Preview
Loading…
Cancel
Save