From bac66127523f640b583b25471284442057f29ca9 Mon Sep 17 00:00:00 2001 From: Ling Wang Date: Fri, 20 Sep 2024 02:18:57 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9E=20fix:=20Fix=20echoback=20and=20ec?= =?UTF-8?q?ho=20filter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.lock | 154 ++++++++++++++++++++++++++++++++++---- Cargo.toml | 2 +- python/tester/__init__.py | 9 +-- src/cli/shell.rs | 93 +++++++++++------------ src/exec/cli_exec.rs | 50 +++++++++++-- src/util/util.rs | 2 + tests/test_shell.py | 11 +++ 7 files changed, 245 insertions(+), 76 deletions(-) create mode 100644 tests/test_shell.py diff --git a/Cargo.lock b/Cargo.lock index 87fbca2..504621f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -367,6 +367,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "downcast-rs" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b325c5dbd37f80359721ad39aca5a29fb04c89279657cffdda8736d0c0b9d2" + [[package]] name = "either" version = "1.13.0" @@ -421,6 +427,17 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "filedescriptor" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7199d965852c3bac31f779ef99cbb4537f80e952e2d6aa0ffeb30cce00f4f46e" +dependencies = [ + "libc", + "thiserror", + "winapi", +] + [[package]] name = "flate2" version = "1.0.33" @@ -628,6 +645,15 @@ dependencies = [ "mach2", ] +[[package]] +name = "ioctl-rs" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7970510895cee30b3e9128319f2cefd4bde883a39f38baa279567ba3a7eb97d" +dependencies = [ + "libc", +] + [[package]] name = "itertools" version = "0.12.1" @@ -812,9 +838,9 @@ dependencies = [ [[package]] name = "memoffset" -version = "0.7.1" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5de893c32cde5f383baa4c04c5d6dbdd735cfd4a794b0debdb2bb1b421da5ff4" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" dependencies = [ "autocfg", ] @@ -859,6 +885,20 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086" +[[package]] +name = "nix" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f346ff70e7dbfd675fe90590b92d59ef2de15a8779ae305ebcbfd3f0caf59be4" +dependencies = [ + "autocfg", + "bitflags 1.3.2", + "cfg-if", + "libc", + "memoffset 0.6.5", + "pin-utils", +] + [[package]] name = "nix" version = "0.26.4" @@ -868,8 +908,6 @@ dependencies = [ "bitflags 1.3.2", "cfg-if", "libc", - "memoffset 0.7.1", - "pin-utils", ] [[package]] @@ -1067,6 +1105,27 @@ version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da544ee218f0d287a911e9c99a39a8c9bc8fcad3cb8db5959940044ecfc67265" +[[package]] +name = "portable-pty" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "806ee80c2a03dbe1a9fb9534f8d19e4c0546b790cde8fd1fea9d6390644cb0be" +dependencies = [ + "anyhow", + "bitflags 1.3.2", + "downcast-rs", + "filedescriptor", + "lazy_static", + "libc", + "log", + "nix 0.25.1", + "serial", + "shared_library", + "shell-words", + "winapi", + "winreg", +] + [[package]] name = "ppv-lite86" version = "0.2.20" @@ -1104,15 +1163,6 @@ dependencies = [ "syn", ] -[[package]] -name = "ptyprocess" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e05aef7befb11a210468a2d77d978dde2c6381a0381e33beb575e91f57fe8cf" -dependencies = [ - "nix 0.26.4", -] - [[package]] name = "pyo3" version = "0.22.3" @@ -1379,6 +1429,48 @@ dependencies = [ "serde", ] +[[package]] +name = "serial" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1237a96570fc377c13baa1b88c7589ab66edced652e43ffb17088f003db3e86" +dependencies = [ + "serial-core", + "serial-unix", + "serial-windows", +] + +[[package]] +name = "serial-core" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f46209b345401737ae2125fe5b19a77acce90cd53e1658cda928e4fe9a64581" +dependencies = [ + "libc", +] + +[[package]] +name = "serial-unix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f03fbca4c9d866e24a459cbca71283f545a37f8e3e002ad8c70593871453cab7" +dependencies = [ + "ioctl-rs", + "libc", + "serial-core", + "termios", +] + +[[package]] +name = "serial-windows" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15c6d3b776267a75d31bbdfd5d36c0ca051251caafc285827052bc53bcdc8162" +dependencies = [ + "libc", + "serial-core", +] + [[package]] name = "serialport" version = "4.5.0" @@ -1397,6 +1489,22 @@ dependencies = [ "winapi", ] +[[package]] +name = "shared_library" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a9e7e0f2bfae24d8a5b5a66c5b257a83c7412304311512a0c054cd5e619da11" +dependencies = [ + "lazy_static", + "libc", +] + +[[package]] +name = "shell-words" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24188a676b6ae68c3b2cb3a01be17fbf7240ce009799bb56d5b1409051e78fde" + [[package]] name = "shlex" version = "1.3.0" @@ -1489,6 +1597,15 @@ version = "0.12.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" +[[package]] +name = "termios" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5d9cf598a6d7ce700a4e6a9199da127e6819a61e64b68609683cc9a01b5683a" +dependencies = [ + "libc", +] + [[package]] name = "tester" version = "0.1.0" @@ -1498,7 +1615,7 @@ dependencies = [ "enigo", "image", "nix 0.29.0", - "ptyprocess", + "portable-pty", "pyo3", "rand", "serde", @@ -2044,6 +2161,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "winreg" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" +dependencies = [ + "winapi", +] + [[package]] name = "xcap" version = "0.0.13" diff --git a/Cargo.toml b/Cargo.toml index 87ab1e0..a1bdafe 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,11 +19,11 @@ ssh2 = "0.9.4" toml = "0.8.19" colored = "2.1.0" nix = { version = "0.29.0", features = ["fs", "process", "signal", "term"] } -ptyprocess = "0.4.1" vte = "0.13.0" image = "0.25.2" xcap = "0.0.13" enigo = "0.2.1" +portable-pty = "0.8.1" [toolchain] channel = "nightly" diff --git a/python/tester/__init__.py b/python/tester/__init__.py index 38103ee..ded356d 100644 --- a/python/tester/__init__.py +++ b/python/tester/__init__.py @@ -1,9 +1,8 @@ from .tester import * +from .hook import GlobalCallHook +from .direct_script import DirectScript +from .tty_abst import TtyAbst -__doc__ = tester.__doc__ +__doc__ = None if hasattr(tester, "__all__"): __all__ = tester.__all__ - -from .hook import GlobalCallHook -from .direct_script import DirectScript -from .tty_abst import TtyAbst \ No newline at end of file diff --git a/src/cli/shell.rs b/src/cli/shell.rs index 1e8f262..41265d0 100644 --- a/src/cli/shell.rs +++ b/src/cli/shell.rs @@ -1,12 +1,10 @@ -use ptyprocess::{stream::Stream, PtyProcess}; -use std::io::BufReader; +use portable_pty::{native_pty_system, Child, CommandBuilder, PtyPair, PtySize}; +use std::io::{BufReader, Read}; use std::ops::DerefMut; -use std::os::fd::AsRawFd; use std::{ any::Any, error::Error, io::Write, - process::Command, sync::{Arc, Mutex}, thread::{sleep, spawn, JoinHandle}, time::Duration, @@ -18,9 +16,11 @@ use crate::{consts::SHELL_DURATION, err, info, log, util::anybase::AnyBase}; use super::tty::Tty; pub struct Shell { - inner: Arc>, buff: Arc>>, - proc: PtyProcess, + pty: PtyPair, + child: Box, + reader: Arc>>, + writer: Arc>>, handle: Option>, stop: Arc>, } @@ -31,34 +31,30 @@ impl Shell { info!("Spawn shell process: {}", shell); - let mut inner = Command::new(shell); - inner.args(["-i"]); - let inner = PtyProcess::spawn(inner); - if let Err(e) = inner { - err!("Failed to spawn shell process. Reason: {}", e); - return Err(Box::new(e)); - } - let proc = inner.unwrap(); - let inner = proc.get_pty_stream()?; + let mut cmd = CommandBuilder::new(shell); + cmd.arg("-i"); + let pty_system = native_pty_system(); + let pty = pty_system.openpty(PtySize::default())?; + + let child = pty.slave.spawn_command(cmd)?; - info!( - "Shell process spawned, got streamed... FD: {:?}", - inner.as_raw_fd() - ); + let reader = pty.master.try_clone_reader()?; - let inner = Arc::new(Mutex::new(inner)); + let writer = pty.master.take_writer()?; let mut res = Shell { - inner, buff: Arc::new(Mutex::new(Vec::new())), - proc, + pty, + child, + reader: Arc::new(Mutex::new(reader)), + writer: Arc::new(Mutex::new(writer)), handle: None, stop: Arc::new(Mutex::new(false)), }; let buff = res.buff.clone(); let stop = res.stop.clone(); - let stream = res.inner.clone(); + let reader = res.reader.clone(); let handle = spawn(move || loop { sleep(Duration::from_millis(SHELL_DURATION)); { @@ -69,9 +65,9 @@ impl Shell { } let mut buf = Vec::new(); { - let mut stream = stream.lock().unwrap(); - let mut reader = BufReader::new(stream.deref_mut()); - let sz = try_read(&mut reader, &mut buf); + let mut reader = reader.lock().unwrap(); + let mut r = BufReader::new(reader.deref_mut()); + let sz = try_read(&mut r, &mut buf); if let Err(e) = sz { err!("Failed to read from shell process. Reason: {}", e); return; @@ -104,7 +100,8 @@ impl Shell { } *stop = true; log!("Try to stop shell process"); - self.proc.exit(false).unwrap(); + writeln!(self.pty.master.take_writer().unwrap(), "exit").unwrap(); + self.child.kill().unwrap(); // if let Some(handle) = self.handle.take() { // handle.join().unwrap(); // self.inner.wait().unwrap(); @@ -131,44 +128,44 @@ impl AnyBase for Shell { impl Tty for Shell { fn read(&mut self) -> Result, Box> { let mut res = Vec::new(); - let mut buff = self.buff.lock().unwrap(); + let buff = self.buff.clone(); + let mut buff = buff.lock().unwrap(); res.extend(buff.iter()); buff.clear(); - if !res.is_empty() { - log!("Shell read: {:?}", String::from_utf8_lossy(&res)); - } Ok(res) } fn read_line(&mut self) -> Result, Box> { let mut res = Vec::new(); + let buff = self.buff.clone(); loop { sleep(Duration::from_millis(SHELL_DURATION)); - let mut buff = self.buff.lock().unwrap(); - if buff.is_empty() { - continue; - } - let mut i = 0; - while i < buff.len() { - res.push(buff[i]); - i += 1; + { + let mut buff = buff.lock().unwrap(); + if buff.is_empty() { + continue; + } + let mut i = 0; + while i < buff.len() { + res.push(buff[i]); + i += 1; + if res.ends_with(&[0x0A]) { + break; + } + } + buff.drain(0..i); if res.ends_with(&[0x0A]) { break; } } - buff.drain(0..i); - if res.ends_with(&[0x0A]) { - break; - } } Ok(res) } fn write(&mut self, data: &[u8]) -> Result<(), Box> { - let mut stream = self.inner.lock().unwrap(); - info!("Shell locked..."); - match stream.write_all(data) { + let writer = self.writer.clone(); + let mut writer = writer.lock().unwrap(); + match writer.write_all(data) { Ok(_) => { - stream.flush().unwrap(); - info!("Shell write: {:?}", String::from_utf8_lossy(data)); + writer.flush().unwrap(); Ok(()) } Err(e) => { diff --git a/src/exec/cli_exec.rs b/src/exec/cli_exec.rs index a0fa13f..064eb54 100644 --- a/src/exec/cli_exec.rs +++ b/src/exec/cli_exec.rs @@ -6,7 +6,10 @@ use std::{ }; use crate::{ - consts::DURATION, err, info, cli::tty::{DynTty, InnerTty, Tty, WrapperTty}, util::{anybase::AnyBase, util::rand_string} + cli::tty::{DynTty, InnerTty, Tty, WrapperTty}, + consts::DURATION, + err, info, + util::{anybase::AnyBase, util::rand_string}, }; use super::cli_api::{CliTestApi, SudoCliTestApi}; @@ -24,6 +27,7 @@ impl CliTester { impl CliTester { fn run_command(&mut self, command: &String) -> Result<(), Box> { info!("Write to shell: {}", command); + sleep(Duration::from_millis(DURATION)); self.inner.write(command.as_bytes()) } } @@ -70,8 +74,21 @@ impl InnerTty for CliTester { } } -impl CliTestApi for CliTester { - fn wait_serial(&mut self, expected: &str, timeout: u32) -> Result> { +impl CliTester { + fn filter_assert_echo(&self, expected: &str, buf: &mut Vec) -> Result<(), Box> { + let expected = "echo ".to_owned() + expected; + let expected = expected.as_bytes(); + for (pos, window) in buf.windows(expected.len()).enumerate() { + if window == expected { + let i = pos + expected.len(); + buf.drain(0..=i); + break; + } + } + Ok(()) + } + + fn do_wait_serial(&mut self, expected: &str, timeout: u32, filter_echo_back: Option<&str>) -> Result> { let begin = Instant::now(); let mut buf = Vec::new(); info!("Waiting for string {{{}}}", expected); @@ -79,6 +96,9 @@ impl CliTestApi for CliTester { sleep(Duration::from_millis(DURATION)); let res = self.inner.read()?; buf.extend_from_slice(&res); + if let Some(filter) = filter_echo_back { + self.filter_assert_echo(filter, &mut buf)?; + } let content = String::from_utf8(buf.clone()).unwrap_or_default(); if content.contains(expected) { info!("Matched string {{{}}}", expected); @@ -97,17 +117,23 @@ impl CliTestApi for CliTester { let res = String::from_utf8(buf)?; Ok(res) } +} + +impl CliTestApi for CliTester { + fn wait_serial(&mut self, expected: &str, timeout: u32) -> Result> { + self.do_wait_serial(expected, timeout, None) + } fn script_run(&mut self, script: &str, timeout: u32) -> Result> { let mut cmd = script.to_owned(); let echo_content_rand = String::from_utf8(rand_string(8)).unwrap(); - cmd += "&& echo "; + cmd += " && echo "; cmd += &echo_content_rand; cmd += " \n"; self.run_command(&cmd)?; - self.wait_serial(&echo_content_rand, timeout) + self.do_wait_serial(&echo_content_rand, timeout, Some(&echo_content_rand)) } fn background_script_run(&mut self, script: &str) -> Result<(), Box> { let mut cmd = script.to_owned(); @@ -175,10 +201,18 @@ impl InnerTty for SudoCliTester { } impl CliTestApi for SudoCliTester { - fn wait_serial(&mut self, expected: &str, timeout: u32) -> Result> { + fn wait_serial( + &mut self, + expected: &str, + timeout: u32, + ) -> Result> { self.inner.wait_serial(expected, timeout) } - fn script_run(&mut self, script: &str, timeout: u32) -> Result> { + fn script_run( + &mut self, + script: &str, + timeout: u32, + ) -> Result> { self.inner.script_run(script, timeout) } fn background_script_run(&mut self, script: &str) -> Result<(), Box> { @@ -197,7 +231,7 @@ impl SudoCliTestApi for SudoCliTester { ) -> Result> { let mut cmd = String::from("sudo "); cmd += script; - cmd += "\n"; + cmd += " "; self.inner.script_run(&cmd, timeout) } } diff --git a/src/util/util.rs b/src/util/util.rs index 8d79f74..1afd1b6 100644 --- a/src/util/util.rs +++ b/src/util/util.rs @@ -5,6 +5,8 @@ use std::{ use rand::{distributions::Alphanumeric, thread_rng, Rng}; +use crate::info; + #[macro_export] macro_rules! todo { () => { diff --git a/tests/test_shell.py b/tests/test_shell.py new file mode 100644 index 0000000..fdba73b --- /dev/null +++ b/tests/test_shell.py @@ -0,0 +1,11 @@ +import tester + +if __name__ == "__main__": + s = tester.Shell("bash") + e = tester.Exec(s) + e.script_run("uname -a") + e.script_run("ls") + e.script_run("uname -a") + e.script_run("ls") + e.script_run("uname -a") + e.script_run("ls")