1 files changed, 321 insertions, 78 deletions

diff --git a/test_util/src/pty.rs b/test_util/src/pty.rs
index f3bb2829f..80d06881e 100644
--- a/test_util/src/pty.rs
+++ b/test_util/src/pty.rs
@@ -1,36 +1,253 @@
 // Copyright 2018-2023 the Deno authors. All rights reserved. MIT license.
 
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::io::Read;
+use std::io::Write;
 use std::path::Path;
+use std::time::Duration;
+use std::time::Instant;
+
+use crate::strip_ansi_codes;
+
+/// Points to know about when writing pty tests:
+///
+/// - Consecutive writes cause issues where you might write while a prompt
+///   is not showing. So when you write, always `.expect(...)` on the output.
+/// - Similar to the last point, using `.expect(...)` can help make the test
+///   more deterministic. If the test is flaky, try adding more `.expect(...)`s
+pub struct Pty {
+  pty: Box<dyn SystemPty>,
+  read_bytes: Vec<u8>,
+  last_index: usize,
+}
+
+impl Pty {
+  pub fn new(
+    program: &Path,
+    args: &[&str],
+    cwd: &Path,
+    env_vars: Option<HashMap<String, String>>,
+  ) -> Self {
+    let pty = create_pty(program, args, cwd, env_vars);
+    let mut pty = Self {
+      pty,
+      read_bytes: Vec::new(),
+      last_index: 0,
+    };
+    if args[0] == "repl" && !args.contains(&"--quiet") {
+      // wait for the repl to start up before writing to it
+      pty.expect("exit using ctrl+d, ctrl+c, or close()");
+    }
+    pty
+  }
+
+  pub fn is_supported() -> bool {
+    let is_mac_or_windows = cfg!(target_os = "macos") || cfg!(windows);
+    if is_mac_or_windows && std::env::var("CI").is_ok() {
+      // the pty tests give a ENOTTY error for Mac and don't really start up
+      // on the windows CI for some reason so ignore them for now
+      eprintln!("Ignoring windows CI.");
+      false
+    } else {
+      true
+    }
+  }
+
+  #[track_caller]
+  pub fn write_raw(&mut self, line: impl AsRef<str>) {
+    let line = if cfg!(windows) {
+      line.as_ref().replace('\n', "\r\n")
+    } else {
+      line.as_ref().to_string()
+    };
+    if let Err(err) = self.pty.write(line.as_bytes()) {
+      panic!("{:#}", err)
+    }
+    self.pty.flush().unwrap();
+  }
+
+  #[track_caller]
+  pub fn write_line(&mut self, line: impl AsRef<str>) {
+    self.write_line_raw(&line);
+
+    // expect what was written to show up in the output
+    // due to "pty echo"
+    for line in line.as_ref().lines() {
+      self.expect(line);
+    }
+  }
+
+  /// Writes a line without checking if it's in the output.
+  #[track_caller]
+  pub fn write_line_raw(&mut self, line: impl AsRef<str>) {
+    self.write_raw(format!("{}\n", line.as_ref()));
+  }
+
+  #[track_caller]
+  pub fn read_until(&mut self, end_text: impl AsRef<str>) -> String {
+    self.read_until_with_advancing(|text| {
+      text
+        .find(end_text.as_ref())
+        .map(|index| index + end_text.as_ref().len())
+    })
+  }
+
+  #[track_caller]
+  pub fn expect(&mut self, text: impl AsRef<str>) {
+    self.read_until(text.as_ref());
+  }
+
+  #[track_caller]
+  pub fn expect_any(&mut self, texts: &[&str]) {
+    self.read_until_with_advancing(|text| {
+      for find_text in texts {
+        if let Some(index) = text.find(find_text) {
+          return Some(index);
+        }
+      }
+      None
+    });
+  }
+
+  /// Consumes and expects to find all the text until a timeout is hit.
+  #[track_caller]
+  pub fn expect_all(&mut self, texts: &[&str]) {
+    let mut pending_texts: HashSet<&&str> = HashSet::from_iter(texts);
+    let mut max_index: Option<usize> = None;
+    self.read_until_with_advancing(|text| {
+      for pending_text in pending_texts.clone() {
+        if let Some(index) = text.find(pending_text) {
+          let index = index + pending_text.len();
+          match &max_index {
+            Some(current) => {
+              if *current < index {
+                max_index = Some(index);
+              }
+            }
+            None => {
+              max_index = Some(index);
+            }
+          }
+          pending_texts.remove(pending_text);
+        }
+      }
+      if pending_texts.is_empty() {
+        max_index
+      } else {
+        None
+      }
+    });
+  }
+
+  /// Expects the raw text to be found, which may include ANSI codes.
+  /// Note: this expects the raw bytes in any output that has already
+  /// occurred or may occur within the next few seconds.
+  #[track_caller]
+  pub fn expect_raw_in_current_output(&mut self, text: impl AsRef<str>) {
+    self.read_until_condition(|pty| {
+      let data = String::from_utf8_lossy(&pty.read_bytes);
+      data.contains(text.as_ref())
+    });
+  }
+
+  #[track_caller]
+  fn read_until_with_advancing(
+    &mut self,
+    mut condition: impl FnMut(&str) -> Option<usize>,
+  ) -> String {
+    let mut final_text = String::new();
+    self.read_until_condition(|pty| {
+      let text = pty.next_text();
+      if let Some(end_index) = condition(&text) {
+        pty.last_index += end_index;
+        final_text = text[..end_index].to_string();
+        true
+      } else {
+        false
+      }
+    });
+    final_text
+  }
+
+  #[track_caller]
+  fn read_until_condition(
+    &mut self,
+    mut condition: impl FnMut(&mut Self) -> bool,
+  ) {
+    let timeout_time =
+      Instant::now().checked_add(Duration::from_secs(5)).unwrap();
+    while Instant::now() < timeout_time {
+      self.fill_more_bytes();
+      if condition(self) {
+        return;
+      }
+    }
 
-pub trait Pty: Read {
-  fn write_text(&mut self, text: &str);
+    let text = self.next_text();
+    eprintln!(
+      "------ Start Full Text ------\n{:?}\n------- End Full Text -------",
+      String::from_utf8_lossy(&self.read_bytes)
+    );
+    eprintln!("Next text: {:?}", text);
+    panic!("Timed out.")
+  }
 
-  fn write_line(&mut self, text: &str) {
-    self.write_text(&format!("{text}\n"));
+  fn next_text(&self) -> String {
+    let text = String::from_utf8_lossy(&self.read_bytes).to_string();
+    let text = strip_ansi_codes(&text);
+    text[self.last_index..].to_string()
   }
 
-  /// Reads the output to the EOF.
-  fn read_all_output(&mut self) -> String {
-    let mut text = String::new();
-    self.read_to_string(&mut text).unwrap();
-    text
+  fn fill_more_bytes(&mut self) {
+    let mut buf = [0; 256];
+    if let Ok(count) = self.pty.read(&mut buf) {
+      self.read_bytes.extend(&buf[..count]);
+    } else {
+      std::thread::sleep(Duration::from_millis(10));
+    }
   }
 }
 
+trait SystemPty: Read + Write {}
+
 #[cfg(unix)]
-pub fn create_pty(
-  program: impl AsRef<Path>,
+fn setup_pty(master: &pty2::fork::Master) {
+  use nix::fcntl::fcntl;
+  use nix::fcntl::FcntlArg;
+  use nix::fcntl::OFlag;
+  use nix::sys::termios;
+  use nix::sys::termios::tcgetattr;
+  use nix::sys::termios::tcsetattr;
+  use nix::sys::termios::SetArg;
+  use std::os::fd::AsRawFd;
+
+  let fd = master.as_raw_fd();
+  let mut term = tcgetattr(fd).unwrap();
+  // disable cooked mode
+  term.local_flags.remove(termios::LocalFlags::ICANON);
+  tcsetattr(fd, SetArg::TCSANOW, &term).unwrap();
+
+  // turn on non-blocking mode so we get timeouts
+  let flags = fcntl(fd, FcntlArg::F_GETFL).unwrap();
+  let new_flags = OFlag::from_bits_truncate(flags) | OFlag::O_NONBLOCK;
+  fcntl(fd, FcntlArg::F_SETFL(new_flags)).unwrap();
+}
+
+#[cfg(unix)]
+fn create_pty(
+  program: &Path,
   args: &[&str],
-  cwd: impl AsRef<Path>,
+  cwd: &Path,
   env_vars: Option<HashMap<String, String>>,
-) -> Box<dyn Pty> {
+) -> Box<dyn SystemPty> {
   let fork = pty2::fork::Fork::from_ptmx().unwrap();
   if fork.is_parent().is_ok() {
+    let master = fork.is_parent().unwrap();
+    setup_pty(&master);
     Box::new(unix::UnixPty { fork })
   } else {
-    std::process::Command::new(program.as_ref())
+    std::process::Command::new(program)
       .current_dir(cwd)
       .args(args)
       .envs(env_vars.unwrap_or_default())
@@ -47,7 +264,7 @@ mod unix {
   use std::io::Read;
   use std::io::Write;
 
-  use super::Pty;
+  use super::SystemPty;
 
   pub struct UnixPty {
     pub fork: pty2::fork::Fork,
@@ -55,46 +272,55 @@ mod unix {
 
   impl Drop for UnixPty {
     fn drop(&mut self) {
-      self.fork.wait().unwrap();
-    }
-  }
+      use nix::sys::signal::kill;
+      use nix::sys::signal::Signal;
+      use nix::unistd::Pid;
 
-  impl Pty for UnixPty {
-    fn write_text(&mut self, text: &str) {
-      let mut master = self.fork.is_parent().unwrap();
-      master.write_all(text.as_bytes()).unwrap();
+      if let pty2::fork::Fork::Parent(child_pid, _) = self.fork {
+        let pid = Pid::from_raw(child_pid);
+        kill(pid, Signal::SIGTERM).unwrap()
+      }
     }
   }
 
+  impl SystemPty for UnixPty {}
+
   impl Read for UnixPty {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
       let mut master = self.fork.is_parent().unwrap();
       master.read(buf)
     }
   }
+
+  impl Write for UnixPty {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+      let mut master = self.fork.is_parent().unwrap();
+      master.write(buf)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+      let mut master = self.fork.is_parent().unwrap();
+      master.flush()
+    }
+  }
 }
 
 #[cfg(target_os = "windows")]
-pub fn create_pty(
-  program: impl AsRef<Path>,
+fn create_pty(
+  program: &Path,
   args: &[&str],
-  cwd: impl AsRef<Path>,
+  cwd: &Path,
   env_vars: Option<HashMap<String, String>>,
-) -> Box<dyn Pty> {
-  let pty = windows::WinPseudoConsole::new(
-    program,
-    args,
-    &cwd.as_ref().to_string_lossy(),
-    env_vars,
-  );
+) -> Box<dyn SystemPty> {
+  let pty = windows::WinPseudoConsole::new(program, args, cwd, env_vars);
   Box::new(pty)
 }
 
 #[cfg(target_os = "windows")]
 mod windows {
   use std::collections::HashMap;
+  use std::io::ErrorKind;
   use std::io::Read;
-  use std::io::Write;
   use std::path::Path;
   use std::ptr;
   use std::time::Duration;
@@ -105,11 +331,13 @@ mod windows {
   use winapi::shared::winerror::S_OK;
   use winapi::um::consoleapi::ClosePseudoConsole;
   use winapi::um::consoleapi::CreatePseudoConsole;
+  use winapi::um::fileapi::FlushFileBuffers;
   use winapi::um::fileapi::ReadFile;
   use winapi::um::fileapi::WriteFile;
   use winapi::um::handleapi::DuplicateHandle;
   use winapi::um::handleapi::INVALID_HANDLE_VALUE;
   use winapi::um::namedpipeapi::CreatePipe;
+  use winapi::um::namedpipeapi::PeekNamedPipe;
   use winapi::um::processthreadsapi::CreateProcessW;
   use winapi::um::processthreadsapi::DeleteProcThreadAttributeList;
   use winapi::um::processthreadsapi::GetCurrentProcess;
@@ -127,7 +355,7 @@ mod windows {
   use winapi::um::winnt::DUPLICATE_SAME_ACCESS;
   use winapi::um::winnt::HANDLE;
 
-  use super::Pty;
+  use super::SystemPty;
 
   macro_rules! assert_win_success {
     ($expression:expr) => {
@@ -138,6 +366,15 @@ mod windows {
     };
   }
 
+  macro_rules! handle_err {
+    ($expression:expr) => {
+      let success = $expression;
+      if success != TRUE {
+        return Err(std::io::Error::last_os_error());
+      }
+    };
+  }
+
   pub struct WinPseudoConsole {
     stdin_write_handle: WinHandle,
     stdout_read_handle: WinHandle,
@@ -149,9 +386,9 @@ mod windows {
 
   impl WinPseudoConsole {
     pub fn new(
-      program: impl AsRef<Path>,
+      program: &Path,
       args: &[&str],
-      cwd: &str,
+      cwd: &Path,
       maybe_env_vars: Option<HashMap<String, String>>,
     ) -> Self {
       // https://docs.microsoft.com/en-us/windows/console/creating-a-pseudoconsole-session
@@ -184,15 +421,19 @@ mod windows {
         let mut proc_info: PROCESS_INFORMATION = std::mem::zeroed();
         let command = format!(
           "\"{}\" {}",
-          program.as_ref().to_string_lossy(),
-          args.join(" ")
+          program.to_string_lossy(),
+          args
+            .iter()
+            .map(|a| format!("\"{}\"", a))
+            .collect::<Vec<_>>()
+            .join(" ")
         )
         .trim()
         .to_string();
-        let mut application_str =
-          to_windows_str(&program.as_ref().to_string_lossy());
+        let mut application_str = to_windows_str(&program.to_string_lossy());
         let mut command_str = to_windows_str(&command);
-        let mut cwd = to_windows_str(cwd);
+        let cwd = cwd.to_string_lossy().replace('/', "\\");
+        let mut cwd = to_windows_str(&cwd);
 
         assert_win_success!(CreateProcessW(
           application_str.as_mut_ptr(),
@@ -242,45 +483,47 @@ mod windows {
 
   impl Read for WinPseudoConsole {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-      loop {
-        let mut bytes_read = 0;
-        // SAFETY:
-        // winapi call
-        let success = unsafe {
-          ReadFile(
-            self.stdout_read_handle.as_raw_handle(),
-            buf.as_mut_ptr() as _,
-            buf.len() as u32,
-            &mut bytes_read,
-            ptr::null_mut(),
-          )
-        };
-
-        // ignore zero-byte writes
-        let is_zero_byte_write = bytes_read == 0 && success == TRUE;
-        if !is_zero_byte_write {
-          return Ok(bytes_read as usize);
-        }
+      // don't do a blocking read in order to support timing out
+      let mut bytes_available = 0;
+      // SAFETY: winapi call
+      handle_err!(unsafe {
+        PeekNamedPipe(
+          self.stdout_read_handle.as_raw_handle(),
+          ptr::null_mut(),
+          0,
+          ptr::null_mut(),
+          &mut bytes_available,
+          ptr::null_mut(),
+        )
+      });
+      if bytes_available == 0 {
+        return Err(std::io::Error::new(ErrorKind::WouldBlock, "Would block."));
       }
-    }
-  }
 
-  impl Pty for WinPseudoConsole {
-    fn write_text(&mut self, text: &str) {
-      // windows pseudo console requires a \r\n to do a newline
-      let newline_re = regex::Regex::new("\r?\n").unwrap();
-      self
-        .write_all(newline_re.replace_all(text, "\r\n").as_bytes())
-        .unwrap();
+      let mut bytes_read = 0;
+      // SAFETY: winapi call
+      handle_err!(unsafe {
+        ReadFile(
+          self.stdout_read_handle.as_raw_handle(),
+          buf.as_mut_ptr() as _,
+          buf.len() as u32,
+          &mut bytes_read,
+          ptr::null_mut(),
+        )
+      });
+
+      Ok(bytes_read as usize)
     }
   }
 
+  impl SystemPty for WinPseudoConsole {}
+
   impl std::io::Write for WinPseudoConsole {
     fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
       let mut bytes_written = 0;
       // SAFETY:
       // winapi call
-      assert_win_success!(unsafe {
+      handle_err!(unsafe {
         WriteFile(
           self.stdin_write_handle.as_raw_handle(),
           buffer.as_ptr() as *const _,
@@ -293,6 +536,10 @@ mod windows {
     }
 
     fn flush(&mut self) -> std::io::Result<()> {
+      // SAFETY: winapi call
+      handle_err!(unsafe {
+        FlushFileBuffers(self.stdin_write_handle.as_raw_handle())
+      });
       Ok(())
     }
   }
@@ -307,12 +554,10 @@ mod windows {
     }
 
     pub fn duplicate(&self) -> WinHandle {
-      // SAFETY:
-      // winapi call
+      // SAFETY: winapi call
       let process_handle = unsafe { GetCurrentProcess() };
       let mut duplicate_handle = ptr::null_mut();
-      // SAFETY:
-      // winapi call
+      // SAFETY: winapi call
       assert_win_success!(unsafe {
         DuplicateHandle(
           process_handle,
@@ -410,8 +655,7 @@ mod windows {
 
   impl Drop for ProcThreadAttributeList {
     fn drop(&mut self) {
-      // SAFETY:
-      // winapi call
+      // SAFETY: winapi call
       unsafe { DeleteProcThreadAttributeList(self.as_mut_ptr()) };
     }
   }
@@ -420,8 +664,7 @@ mod windows {
     let mut read_handle = std::ptr::null_mut();
     let mut write_handle = std::ptr::null_mut();
 
-    // SAFETY:
-    // Creating an anonymous pipe with winapi.
+    // SAFETY: Creating an anonymous pipe with winapi.
     assert_win_success!(unsafe {
       CreatePipe(&mut read_handle, &mut write_handle, ptr::null_mut(), 0)
     });