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,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() } {
+fn spawn_child(ptys: &Ptys) {
-        Ok(ForkResult::Child) => {
     let cmd = CString::new("/bin/bash").unwrap();
     let bcmd = std::env::args().skip(1).collect::<Vec<String>>().join(" ");
     println!("Cmd: '{}'", bcmd);
@@ -117,10 +139,9 @@ fn main() {
         CString::new(bcmd).unwrap(),
     ];
     execvp(&cmd, &args).expect("execvp failed");
-        }
+}
 
-        Ok(ForkResult::Parent { child }) => {
+fn main_loop(child: Pid, ptys: &Ptys) {
-            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.
@@ -162,40 +183,25 @@ fn main() {
             std::process::exit(-1);
         }
 
-                let revents = poll_fds[0].revents();
+        match handle_poll(
-                if let Some(revents) = revents {
+            &poll_fds[1],
-                    if revents.contains(PollFlags::POLLHUP) {
+            StdioStreamType::StdErr,
-                        break;
+            &master_file_err,
-                    }
+            &mut err_buf,
-                    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,
+            Err(true) => break,
-                        }
+            Err(false) => continue,
-                    }
+        }
-                }
+        match handle_poll(
-
+            &poll_fds[0],
-                let revents = poll_fds[1].revents();
+            StdioStreamType::StdOut,
-                if let Some(revents) = revents {
+            &master_file_out,
-                    if revents.contains(PollFlags::POLLHUP) {
+            &mut out_buf,
-                        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,
+            Err(true) => break,
-                        }
+            Err(false) => continue,
-                    }
         }
 
         if res.ok().unwrap() == 0 && unsafe { kill(child.as_raw(), 0) } != 0 {
@@ -214,6 +220,19 @@ fn main() {
 
         waitpid(child, Some(WaitPidFlag::WNOHANG)).ok();
     }
+}
+
+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) => {
@@ -221,5 +240,4 @@ fn main() {
             std::process::exit(-1);
         }
     }
-    println!("");
 }