Minor cleanup

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

@ -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,15 +82,38 @@ 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) => {
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);
@ -119,8 +141,7 @@ fn main() {
execvp(&cmd, &args).expect("execvp failed");
}
Ok(ForkResult::Parent { child }) => {
let _ = child;
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.
@ -162,40 +183,25 @@ fn main() {
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,
}) {
match handle_poll(
&poll_fds[1],
StdioStreamType::StdErr,
&master_file_err,
&mut err_buf,
) {
Ok(_) => {}
Err(_) => break,
}
}
}
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,
}) {
Err(true) => break,
Err(false) => continue,
}
match handle_poll(
&poll_fds[0],
StdioStreamType::StdOut,
&master_file_out,
&mut out_buf,
) {
Ok(_) => {}
Err(_) => break,
}
}
Err(true) => break,
Err(false) => continue,
}
if res.ok().unwrap() == 0 && unsafe { kill(child.as_raw(), 0) } != 0 {
@ -216,10 +222,22 @@ fn main() {
}
}
fn main() {
let mut ptys = Ptys::default();
open_ptys(&mut ptys);
match unsafe { fork() } {
Ok(ForkResult::Child) => {
spawn_child(&ptys);
}
Ok(ForkResult::Parent { child }) => {
main_loop(child, &ptys);
}
Err(e) => {
eprintln!("Fork failed: {}", e);
std::process::exit(-1);
}
}
println!("");
}

Loading…
Cancel
Save