Minor cleanup

Deduplicated some code into functions. Behavior should be equivalent to
the behavior of the last commit.

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<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() {
+fn open_ptys(ptys: &mut Ptys) {
-    let mut ptys = Ptys::default();
     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() } {
+fn spawn_child(ptys: &Ptys) {
-        Ok(ForkResult::Child) => {
+    let cmd = CString::new("/bin/bash").unwrap();
-            let cmd = CString::new("/bin/bash").unwrap();
+    let bcmd = std::env::args().skip(1).collect::<Vec<String>>().join(" ");
-            let bcmd = std::env::args().skip(1).collect::<Vec<String>>().join(" ");
+    println!("Cmd: '{}'", bcmd);
-            println!("Cmd: '{}'", bcmd);
+
-
+    setsid().expect("Failed to create new session");
-            setsid().expect("Failed to create new session");
+
-
+    unsafe { libc::close(ptys.out_master.as_raw_fd()) };
-            unsafe { libc::close(ptys.out_master.as_raw_fd()) };
+    unsafe {
-            unsafe {
+        libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDOUT_FILENO);
-                libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDOUT_FILENO);
+        libc::dup2(ptys.err_slave.as_raw_fd(), libc::STDERR_FILENO);
-                libc::dup2(ptys.err_slave.as_raw_fd(), libc::STDERR_FILENO);
+        // There's no real reason to dup the stdin.
-                // There's no real reason to dup the stdin.
+        // It's not being used anyway.
-                // It's not being used anyway.
+        // ----------------------------------------------------------
-                // ----------------------------------------------------------
+        // libc::dup2(ptys.out_slave.as_raw_fd(), libc::STDIN_FILENO);
-                // 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()) };
+fn main_loop(child: Pid, ptys: &Ptys) {
-            unsafe { libc::close(ptys.err_slave.as_raw_fd()) };
+    // 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 = [
+        match handle_poll(
-                cmd.clone(),
+            &poll_fds[1],
-                CString::new("-c").unwrap(),
+            StdioStreamType::StdErr,
-                CString::new(bcmd).unwrap(),
+            &master_file_err,
-            ];
+            &mut err_buf,
-            execvp(&cmd, &args).expect("execvp failed");
+        ) {
+            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 }) => {
+        if res.ok().unwrap() == 0 && unsafe { kill(child.as_raw(), 0) } != 0 {
-            let _ = child;
+            // The poll might've timed out and in the short time between
-            // TODO: The kernel seems to garbage collect the pty
+            // the poll timeout and the check if the process is alive,
-            // & the associated data when the only holder of the slave end dies,
+            // there might be pending data (unless the kernel deletes it...)
-            // resulting in errors reading it. So we keep a copy around.
+            // I have no evidence supporting this theory, though.
-            // I don't actually know if this is the correct way of doing things.
+            // -------------------------------------------------------------
-            // But it fixed the issues, so...
+            // let mut buf = [0u8; 64];
-            // -----------------------------------------------------------------
+            // match master_file.read(&mut buf) {
-            // unsafe { libc::close(slave.as_raw_fd()) };
+            //     Ok(e) => println!("Ok: {}", e),
-
+            //     Err(_e) => {/*println!("Err: {}", e)*/}
-            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);
             // }
+            break;
+        }
 
-            let mut poll_fds = [
+        waitpid(child, Some(WaitPidFlag::WNOHANG)).ok();
-                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,
-                        }
-                    }
-                }
 
-                let revents = poll_fds[1].revents();
+fn main() {
-                if let Some(revents) = revents {
+    let mut ptys = Ptys::default();
-                    if revents.contains(PollFlags::POLLHUP) {
+    open_ptys(&mut ptys);
-                        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,
-                        }
-                    }
-                }
 
-                if res.ok().unwrap() == 0 && unsafe { kill(child.as_raw(), 0) } != 0 {
+    match unsafe { fork() } {
-                    // The poll might've timed out and in the short time between
+        Ok(ForkResult::Child) => {
-                    // the poll timeout and the check if the process is alive,
+            spawn_child(&ptys);
-                    // 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();
+        Ok(ForkResult::Parent { child }) => {
-            }
+            main_loop(child, &ptys);
         }
 
         Err(e) => {
@@ -221,5 +240,4 @@ fn main() {
             std::process::exit(-1);
         }
     }
-    println!("");
 }