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

Add a kernel order mode and plain mode

Browse Source 
The kernel order flag is a workaround for Issue #1. Plain mode simply
removes all of the fancyness that's normally printed.
master
itycodes 3 weeks ago
parent 20ffb15281
commit e5fc3b149e
1 changed files with 172 additions and 56 deletions
Show Stats Download Patch File Download Diff File

228
src/main.rs
Unescape View File

@ -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<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) {
@ -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::<Vec<String>>().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<u8> = Vec::new();
let mut err_buf: Vec<u8> = Vec::new();
let mut out_buf: Vec<u8> = Vec::with_capacity(65536);
let mut err_buf: Vec<u8> = 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::<Vec<OsString>>();
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<OsString> = 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::<Vec<_>>()
.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 }) => {

Write Preview
Loading…
Cancel
Save