├── src ├── io │ ├── util │ │ ├── mod.rs │ │ ├── cache │ │ │ ├── mod.rs │ │ │ ├── async_read_cache_ext.rs │ │ │ ├── read_until_regex.rs │ │ │ └── read_until.rs │ │ └── timeout │ │ │ ├── mod.rs │ │ │ ├── async_read_timeout_ext.rs │ │ │ ├── async_read_cache_timeout_ext.rs │ │ │ ├── read_until_timeout.rs │ │ │ ├── read_until_regex_timeout.rs │ │ │ ├── read_to_end_timeout.rs │ │ │ ├── read_timeout.rs │ │ │ └── read_exact_timeout.rs │ ├── pipe │ │ ├── convert │ │ │ ├── mod.rs │ │ │ ├── stream.rs │ │ │ ├── tcp.rs │ │ │ ├── conv_macro.rs │ │ │ └── process.rs │ │ ├── interactive.rs │ │ ├── error.rs │ │ ├── write.rs │ │ ├── mod.rs │ │ ├── ansi.rs │ │ └── read.rs │ ├── payload │ │ ├── builder │ │ │ ├── build.rs │ │ │ ├── arch.rs │ │ │ ├── condition.rs │ │ │ ├── dynamic_payload.rs │ │ │ ├── convert.rs │ │ │ ├── mod.rs │ │ │ ├── print.rs │ │ │ ├── send.rs │ │ │ └── read.rs │ │ ├── payloads │ │ │ ├── mod.rs │ │ │ ├── initial.rs │ │ │ ├── condition.rs │ │ │ ├── convert.rs │ │ │ ├── try_convert.rs │ │ │ ├── chain.rs │ │ │ ├── dynamic_payload.rs │ │ │ ├── send.rs │ │ │ ├── read.rs │ │ │ └── print.rs │ │ └── mod.rs │ ├── mod.rs │ ├── cache │ │ ├── traits.rs │ │ ├── mod.rs │ │ └── read.rs │ ├── merge │ │ └── mod.rs │ └── stdio │ │ ├── mod.rs │ │ ├── ncurses_bridge.rs │ │ └── shell_bridge.rs ├── lib.rs └── unix │ ├── mod.rs │ ├── error.rs │ ├── symbol.rs │ ├── got.rs │ ├── elf_file.rs │ └── plt.rs ├── .gitignore ├── Cargo.toml ├── LICENSE └── README.md /src/io/util/mod.rs: -------------------------------------------------------------------------------- 1 | mod cache; 2 | mod timeout; 3 | 4 | pub use cache::*; 5 | pub use timeout::*; 6 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | #[allow(async_fn_in_trait)] 2 | #[allow(unused_imports)] 3 | pub mod io; 4 | pub mod unix; 5 | -------------------------------------------------------------------------------- /src/io/pipe/convert/mod.rs: -------------------------------------------------------------------------------- 1 | mod conv_macro; 2 | mod process; 3 | mod stream; 4 | mod tcp; 5 | 6 | pub use conv_macro::*; 7 | pub use process::*; 8 | pub use stream::*; 9 | pub use tcp::*; 10 | -------------------------------------------------------------------------------- /src/io/util/cache/mod.rs: -------------------------------------------------------------------------------- 1 | mod async_read_cache_ext; 2 | mod read_until; 3 | mod read_until_regex; 4 | 5 | pub use async_read_cache_ext::*; 6 | pub use read_until::*; 7 | pub use read_until_regex::*; 8 | -------------------------------------------------------------------------------- /src/unix/mod.rs: -------------------------------------------------------------------------------- 1 | mod elf_file; 2 | mod error; 3 | mod got; 4 | mod plt; 5 | mod symbol; 6 | 7 | pub use elf_file::*; 8 | pub use error::*; 9 | pub use symbol::*; 10 | pub(crate) use got::*; 11 | pub(crate) use plt::*; 12 | -------------------------------------------------------------------------------- /src/io/payload/builder/build.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::*; 3 | 4 | impl PayloadBuilder { 5 | pub fn build(self) -> T { 6 | self.payload 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /src/io/mod.rs: -------------------------------------------------------------------------------- 1 | mod payload; 2 | mod pipe; 3 | 4 | pub mod cache; 5 | mod merge; 6 | pub mod stdio; 7 | mod util; 8 | 9 | pub(crate) use cache::*; 10 | pub use payload::*; 11 | pub use pipe::*; 12 | pub(crate) use stdio::*; 13 | pub use util::*; 14 | -------------------------------------------------------------------------------- /src/unix/error.rs: -------------------------------------------------------------------------------- 1 | use thiserror::Error; 2 | 3 | #[derive(Error, Debug)] 4 | pub enum ElfError { 5 | #[error("io error {0}")] 6 | IOError(#[from] std::io::Error), 7 | #[error("elf parse error {0}")] 8 | ParseError(#[from] elf::ParseError), 9 | #[error("unknown error")] 10 | Unknown, 11 | } 12 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Generated by Cargo 2 | # will have compiled files and executables 3 | /target/ 4 | 5 | # Remove Cargo.lock from gitignore if creating an executable, leave it for libraries 6 | # More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html 7 | Cargo.lock 8 | 9 | # These are backup files generated by rustfmt 10 | **/*.rs.bk 11 | 12 | .idea/* -------------------------------------------------------------------------------- /src/io/pipe/convert/stream.rs: -------------------------------------------------------------------------------- 1 | use crate::io::Pipe; 2 | use tokio::io::{split, ReadHalf, WriteHalf}; 3 | use tokio::io::{AsyncRead, AsyncWrite}; 4 | 5 | impl From for Pipe, WriteHalf> { 6 | fn from(value: T) -> Self { 7 | let (read_stream, write_stream) = split(value); 8 | Pipe::new(read_stream, write_stream) 9 | } 10 | } 11 | -------------------------------------------------------------------------------- /src/io/payload/builder/arch.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::*; 3 | 4 | impl PayloadBuilder { 5 | pub fn x86(self) -> PayloadBuilder { 6 | PayloadBuilder::from(self.payload) 7 | } 8 | 9 | pub fn x64(self) -> PayloadBuilder { 10 | PayloadBuilder::from(self.payload) 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /src/io/payload/payloads/mod.rs: -------------------------------------------------------------------------------- 1 | mod chain; 2 | mod condition; 3 | mod convert; 4 | mod dynamic_payload; 5 | mod initial; 6 | mod print; 7 | mod read; 8 | mod send; 9 | mod try_convert; 10 | 11 | pub use chain::*; 12 | pub use condition::*; 13 | pub use convert::*; 14 | pub use dynamic_payload::*; 15 | pub use initial::*; 16 | pub use print::*; 17 | pub use read::*; 18 | pub use send::*; 19 | pub use try_convert::*; 20 | -------------------------------------------------------------------------------- /src/io/payload/builder/condition.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::{Condition, PayloadAction, ReturnsValue}; 3 | 4 | impl PayloadBuilder 5 | { 6 | pub fn condition( 7 | self, 8 | action: F, 9 | ) -> PayloadBuilder, A> 10 | where F: Fn(&T::ReturnType) -> bool + Clone 11 | { 12 | PayloadBuilder::from(Condition::::new(self.payload, action)) 13 | } 14 | } 15 | -------------------------------------------------------------------------------- /src/io/cache/traits.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::pin::Pin; 3 | use std::task::{Context, Poll}; 4 | pub use tokio::io::AsyncBufRead; 5 | use tokio::io::{AsyncRead, ReadBuf}; 6 | 7 | pub trait AsyncCacheRead: AsyncRead { 8 | fn poll_reader( 9 | self: Pin<&mut Self>, 10 | cx: &mut Context<'_>, 11 | buf: &mut ReadBuf<'_>, 12 | ) -> Poll>; 13 | 14 | fn consume(self: Pin<&mut Self>, amt: usize); 15 | fn restore(self: Pin<&mut Self>, data: &[u8]); 16 | } 17 | -------------------------------------------------------------------------------- /src/io/payload/builder/dynamic_payload.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::payload::payloads::DynamicPayload; 3 | use crate::io::*; 4 | 5 | impl PayloadBuilder { 6 | pub fn payload( 7 | self, 8 | action: F, 9 | ) -> PayloadBuilder, A> 10 | where 11 | F: Fn(T::ReturnType) -> R + Clone, 12 | R: PayloadAction, 13 | { 14 | PayloadBuilder::from(DynamicPayload::::new(self.payload, action)) 15 | } 16 | } 17 | -------------------------------------------------------------------------------- /src/io/payload/payloads/initial.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | 3 | #[derive(Clone)] 4 | pub struct Initial {} 5 | 6 | impl Sendable for Initial {} 7 | impl Readable for Initial {} 8 | 9 | impl Initial { 10 | pub fn new() -> Initial { 11 | Initial {} 12 | } 13 | } 14 | 15 | impl PayloadAction for Initial { 16 | type ReturnType = (); 17 | 18 | async fn execute( 19 | &self, 20 | _pipe: &mut T, 21 | ) -> Result { 22 | Ok(()) 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /src/io/cache/mod.rs: -------------------------------------------------------------------------------- 1 | mod read; 2 | mod traits; 3 | 4 | pub use traits::*; 5 | 6 | use pin_project_lite::pin_project; 7 | 8 | pin_project! { 9 | /// An `AsyncRead`er which applies a timeout to read operations. 10 | #[derive(Debug)] 11 | pub struct CacheReader { 12 | #[pin] 13 | pub(crate) reader: R, 14 | #[pin] 15 | pub(crate) cache: Vec 16 | } 17 | } 18 | 19 | impl CacheReader { 20 | pub fn new(reader: R) -> CacheReader { 21 | CacheReader { 22 | reader, 23 | cache: Vec::new(), 24 | } 25 | } 26 | } 27 | -------------------------------------------------------------------------------- /src/io/pipe/convert/tcp.rs: -------------------------------------------------------------------------------- 1 | use crate::io::Pipe; 2 | use tokio::{ 3 | io::Result, 4 | net::{tcp::*, *}, 5 | }; 6 | 7 | pub type TcpPipe = Pipe; 8 | 9 | impl TcpPipe { 10 | pub async fn connect(addr: A) -> Result { 11 | let stream = TcpStream::connect(addr).await?; 12 | Ok(stream.into()) 13 | } 14 | } 15 | 16 | impl From for TcpPipe { 17 | fn from(value: TcpStream) -> Self { 18 | let (read_stream, write_stream) = value.into_split(); 19 | Pipe::new(read_stream, write_stream) 20 | } 21 | } 22 | -------------------------------------------------------------------------------- /src/io/pipe/interactive.rs: -------------------------------------------------------------------------------- 1 | use crate::io::{NcursesTerminalBridge, PipeError, ShellTerminalBridge, TerminalBridge}; 2 | use tokio::io::{split, AsyncRead, AsyncWrite}; 3 | 4 | impl PipeInteractiveExt for T {} 5 | 6 | pub trait PipeInteractiveExt: AsyncRead + AsyncWrite + Unpin + Send { 7 | async fn interactive_shell(&mut self) -> Result<(), PipeError> { 8 | let (mut reader, mut writer) = split(self); 9 | ShellTerminalBridge::bridge(&mut reader, &mut writer).await; 10 | Ok(()) 11 | } 12 | 13 | async fn interactive_ansi(&mut self) -> Result<(), PipeError> { 14 | let (mut reader, mut writer) = split(self); 15 | NcursesTerminalBridge::bridge(&mut reader, &mut writer).await; 16 | Ok(()) 17 | } 18 | } 19 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "ctf-pwn" 3 | version = "0.3.8" 4 | edition = "2021" 5 | authors = ["Jovan Pavlovic"] 6 | description = "Pwn utilities for Rust." 7 | license = "MIT" 8 | homepage = "https://github.com/rust-ctf/ctf-pwn" 9 | repository = "https://github.com/rust-ctf/ctf-pwn" 10 | readme = "README.md" 11 | 12 | # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html 13 | 14 | [dependencies] 15 | ascii = "1.1.0" 16 | crossterm = { version = "0.27.0", features = ["bracketed-paste"] } 17 | kmp = "0.1.1" 18 | paste = "1.0.14" 19 | pin-project-lite = "0.2.13" 20 | pin-utils = "0.1.0" 21 | regex = { version = "1.10.2", features = [] } 22 | thiserror = "1.0.50" 23 | tokio = { version = "1.34.0", features = ["full"] } 24 | log = "0.4.20" 25 | hex = "0.4.3" 26 | elf = "0.7.4" 27 | -------------------------------------------------------------------------------- /src/io/payload/builder/convert.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::{Convert, PayloadAction, PipeError, ReturnsValue, TryConvert}; 3 | 4 | impl PayloadBuilder { 5 | pub fn convert( 6 | self, 7 | action: F, 8 | ) -> PayloadBuilder, A> 9 | where F: Fn(T::ReturnType) -> E + Clone, E:Clone 10 | { 11 | PayloadBuilder::from(Convert::::new(self.payload, action)) 12 | } 13 | 14 | pub fn try_convert( 15 | self, 16 | action: F, 17 | ) -> PayloadBuilder, A> 18 | where F: Fn(T::ReturnType) -> Result + Clone, E:Clone 19 | { 20 | PayloadBuilder::from(TryConvert::::new(self.payload, action)) 21 | } 22 | } 23 | -------------------------------------------------------------------------------- /src/io/pipe/error.rs: -------------------------------------------------------------------------------- 1 | use ascii::FromAsciiError; 2 | 3 | use thiserror::Error; 4 | 5 | #[derive(Error, Debug)] 6 | pub enum PipeError { 7 | #[error("io error {0}")] 8 | IOError(#[from] std::io::Error), 9 | #[error("ascii parse error {0}")] 10 | AsciiParseError(String), 11 | #[error("utf8 parse error {0}")] 12 | Utf8ParseError(#[from] std::string::FromUtf8Error), 13 | #[error("format error {0}")] 14 | FmtError(#[from] std::fmt::Error), 15 | #[error("regex error {0}")] 16 | RegexError(#[from] regex::Error), 17 | #[error("recv timeout")] 18 | Timeout, 19 | #[error("recv timeout")] 20 | ConditionFailed, 21 | #[error("unknown error")] 22 | Unknown, 23 | } 24 | 25 | impl From> for PipeError { 26 | fn from(value: FromAsciiError) -> Self { 27 | PipeError::AsciiParseError(format!("{value}")) 28 | } 29 | } 30 | -------------------------------------------------------------------------------- /src/io/payload/mod.rs: -------------------------------------------------------------------------------- 1 | mod builder; 2 | mod payloads; 3 | 4 | pub(crate) use builder::*; 5 | pub(crate) use payloads::*; 6 | 7 | use crate::io::payload::builder::PayloadBuilder; 8 | use crate::io::payload::payloads::Initial; 9 | use crate::io::{AsyncCacheRead, PipeError, PipeRead, PipeReadExt, PipeWrite, PipeWriteExt}; 10 | 11 | pub trait PayloadAction: Clone { 12 | type ReturnType: Clone; 13 | 14 | async fn execute( 15 | &self, 16 | pipe: &mut T, 17 | ) -> Result; 18 | } 19 | 20 | #[derive(Clone)] 21 | pub struct Payload {} 22 | 23 | #[derive(Clone)] 24 | pub struct UnknownArch; 25 | #[derive(Clone)] 26 | pub struct X86; 27 | #[derive(Clone)] 28 | pub struct X64; 29 | 30 | impl Payload { 31 | pub fn builder() -> PayloadBuilder { 32 | PayloadBuilder::::new() 33 | } 34 | } 35 | -------------------------------------------------------------------------------- /src/io/pipe/convert/conv_macro.rs: -------------------------------------------------------------------------------- 1 | #[macro_export] 2 | macro_rules! pipe { 3 | ($stream:expr) => {{ 4 | let pipe: Pipe<_, _> = $stream.into(); 5 | pipe 6 | }}; 7 | } 8 | 9 | #[macro_export] 10 | macro_rules! tcp_pipe { 11 | ($input:expr) => {{ 12 | let stream = TcpStream::connect($input).await.unwrap(); 13 | let pipe: TcpStreamPipe = stream.into(); 14 | pipe 15 | }}; 16 | } 17 | 18 | #[macro_export] 19 | macro_rules! shell_pipe { 20 | ($program:expr $(, $arg:expr)*) => {{ 21 | let mut command = Command::new($program); 22 | $( 23 | command.arg($arg); 24 | )* 25 | 26 | ProcessPipe::spawn_command(command).unwrap() 27 | }}; 28 | } 29 | 30 | #[macro_export] 31 | macro_rules! tcp_shell_pipe { 32 | (true, $input:expr, $(, $arg:expr)*) => {{ 33 | tcp_pipe!($input) 34 | }}; 35 | (false, $program:expr, $(, $arg:expr)*) => {{ 36 | file_pipe!($program $(, $arg)*) 37 | }}; 38 | } 39 | -------------------------------------------------------------------------------- /src/io/merge/mod.rs: -------------------------------------------------------------------------------- 1 | use pin_project_lite::pin_project; 2 | use std::task::Poll; 3 | use tokio::io::*; 4 | 5 | //TODO: AutoGenerate MergedOutput for up to ten args 6 | 7 | pin_project! { 8 | #[derive(Debug)] 9 | pub struct MergedAsyncReader { 10 | #[pin] 11 | pub stream0: R1, 12 | #[pin] 13 | pub stream1: R2, 14 | } 15 | } 16 | 17 | impl MergedAsyncReader { 18 | pub fn new(stream0: R1, stream1: R2) -> MergedAsyncReader { 19 | MergedAsyncReader { stream0, stream1 } 20 | } 21 | } 22 | 23 | impl AsyncRead for MergedAsyncReader { 24 | fn poll_read( 25 | self: std::pin::Pin<&mut Self>, 26 | cx: &mut std::task::Context<'_>, 27 | buf: &mut ReadBuf<'_>, 28 | ) -> Poll> { 29 | let me = self.project(); 30 | match me.stream0.poll_read(cx, buf) { 31 | Poll::Ready(r) => Poll::Ready(r), 32 | Poll::Pending => me.stream1.poll_read(cx, buf), 33 | } 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 rust-ctf 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/io/util/cache/async_read_cache_ext.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::cache::read_until::{read_until, ReadUntil}; 2 | use crate::io::{read_until_regex, AsyncCacheRead, ReadUntilRegex}; 3 | 4 | pub trait AsyncCacheReadExt: AsyncCacheRead { 5 | fn read_until<'a, T: AsRef<[u8]>>( 6 | &'a mut self, 7 | delim: T, 8 | buf: &'a mut Vec, 9 | ) -> ReadUntil<'a, Self, T> 10 | where 11 | Self: Unpin, 12 | { 13 | read_until(self, delim, buf) 14 | } 15 | 16 | fn read_line<'a>(&'a mut self, buf: &'a mut Vec) -> ReadUntil<'a, Self, &'static [u8]> 17 | where 18 | Self: Unpin, 19 | { 20 | read_until(self, b"\n", buf) 21 | } 22 | 23 | fn read_line_crlf<'a>(&'a mut self, buf: &'a mut Vec) -> ReadUntil<'a, Self, &'static [u8]> 24 | where 25 | Self: Unpin, 26 | { 27 | read_until(self, b"\r\n", buf) 28 | } 29 | 30 | async fn read_until_regex<'a>( 31 | &'a mut self, 32 | pattern: &str, 33 | buf: &'a mut Vec, 34 | ) -> Result, regex::Error> 35 | where 36 | Self: Unpin, 37 | { 38 | read_until_regex(self, pattern, buf) 39 | } 40 | } 41 | 42 | impl AsyncCacheReadExt for R {} 43 | -------------------------------------------------------------------------------- /src/io/util/timeout/mod.rs: -------------------------------------------------------------------------------- 1 | mod async_read_cache_timeout_ext; 2 | mod async_read_timeout_ext; 3 | mod read_exact_timeout; 4 | mod read_timeout; 5 | mod read_to_end_timeout; 6 | mod read_until_regex_timeout; 7 | mod read_until_timeout; 8 | 9 | pub use async_read_cache_timeout_ext::*; 10 | pub use async_read_timeout_ext::*; 11 | pub use read_exact_timeout::*; 12 | pub use read_timeout::*; 13 | pub use read_until_regex_timeout::*; 14 | pub use read_until_timeout::*; 15 | use std::io; 16 | use std::io::ErrorKind; 17 | 18 | use std::time::Duration; 19 | use tokio::time::Instant; 20 | 21 | pub(crate) fn get_deadline(timeout: Duration) -> Instant { 22 | Instant::now() 23 | .checked_add(timeout) 24 | .unwrap_or_else(|| far_future()) 25 | } 26 | 27 | pub(crate) fn far_future() -> Instant { 28 | // Roughly 30 years from now. 29 | // API does not provide a way to obtain max `Instant` 30 | // or convert specific date in the future to instant. 31 | // 1000 years overflows on macOS, 100 years overflows on FreeBSD. 32 | Instant::now() + Duration::from_secs(86400 * 365 * 30) 33 | } 34 | 35 | fn eof() -> io::Error { 36 | io::Error::new(ErrorKind::UnexpectedEof, "EOF") 37 | } 38 | fn timeout() -> io::Error { 39 | io::Error::new(ErrorKind::TimedOut, "Timeout") 40 | } 41 | -------------------------------------------------------------------------------- /src/io/payload/payloads/condition.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | 3 | #[derive(Clone)] 4 | pub struct Condition { 5 | prev_payload: P, 6 | action: F, 7 | } 8 | 9 | impl Buildable for Condition where Self: PayloadAction {} 10 | impl Sendable for Condition where Self: PayloadAction {} 11 | impl Readable for Condition where Self: PayloadAction {} 12 | impl ReturnsValue for Condition where Self: PayloadAction {} 13 | 14 | impl Condition 15 | where 16 | P: PayloadAction, 17 | F: Fn(&P::ReturnType) -> bool + Clone 18 | { 19 | pub fn new(prev_payload: P, action: F) -> Condition 20 | { 21 | Condition { 22 | prev_payload, 23 | action, 24 | } 25 | } 26 | } 27 | 28 | impl PayloadAction for Condition 29 | where 30 | P: PayloadAction, 31 | F: Fn(&P::ReturnType) -> bool + Clone 32 | { 33 | type ReturnType = P::ReturnType; 34 | 35 | async fn execute( 36 | &self, 37 | pipe: &mut T1, 38 | ) -> Result { 39 | let result = self.prev_payload.execute(pipe).await?; 40 | if !(self.action)(&result) { 41 | return Err(PipeError::ConditionFailed); 42 | } 43 | Ok(result) 44 | } 45 | } 46 | -------------------------------------------------------------------------------- /src/io/payload/payloads/convert.rs: -------------------------------------------------------------------------------- 1 | use std::marker::PhantomData; 2 | use crate::io::*; 3 | 4 | #[derive(Clone)] 5 | pub struct Convert { 6 | prev_payload: P, 7 | action: F, 8 | _phantom: PhantomData, 9 | } 10 | 11 | impl Buildable for Convert where Self: PayloadAction {} 12 | impl Sendable for Convert where Self: PayloadAction {} 13 | impl Readable for Convert where Self: PayloadAction {} 14 | impl ReturnsValue for Convert where Self: PayloadAction {} 15 | 16 | impl Convert 17 | where 18 | P: PayloadAction, 19 | F: Fn(P::ReturnType) -> E + Clone 20 | { 21 | pub fn new(prev_payload: P, action: F) -> Convert 22 | { 23 | Convert { 24 | prev_payload, 25 | action, 26 | _phantom: PhantomData::default(), 27 | } 28 | } 29 | } 30 | 31 | impl PayloadAction for Convert 32 | where 33 | P: PayloadAction, 34 | E: Clone, 35 | F: Fn(P::ReturnType) -> E + Clone 36 | { 37 | type ReturnType = E; 38 | 39 | async fn execute( 40 | &self, 41 | pipe: &mut T1, 42 | ) -> Result { 43 | let result = self.prev_payload.execute(pipe).await?; 44 | Ok((self.action)(result)) 45 | } 46 | } 47 | -------------------------------------------------------------------------------- /src/io/cache/read.rs: -------------------------------------------------------------------------------- 1 | use crate::io::cache::CacheReader; 2 | use crate::io::AsyncCacheRead; 3 | 4 | use std::borrow::BorrowMut; 5 | use std::pin::Pin; 6 | use std::task::{Context, Poll}; 7 | use tokio::io; 8 | use tokio::io::{AsyncRead, ReadBuf}; 9 | 10 | impl AsyncRead for CacheReader { 11 | fn poll_read( 12 | self: Pin<&mut Self>, 13 | cx: &mut Context<'_>, 14 | buf: &mut ReadBuf<'_>, 15 | ) -> Poll> { 16 | let mut this = self.project(); 17 | if !this.cache.is_empty() { 18 | let remaining = usize::min(buf.remaining(), this.cache.len()); 19 | buf.put_slice(&this.cache[..remaining]); 20 | this.cache.drain(..remaining); 21 | return Poll::Ready(Ok(())); 22 | } 23 | this.reader.poll_read(cx, buf) 24 | } 25 | } 26 | 27 | impl AsyncCacheRead for CacheReader { 28 | fn poll_reader( 29 | self: Pin<&mut Self>, 30 | cx: &mut Context<'_>, 31 | buf: &mut ReadBuf<'_>, 32 | ) -> Poll> { 33 | self.project().reader.poll_read(cx, buf) 34 | } 35 | 36 | fn consume(self: Pin<&mut Self>, amt: usize) { 37 | self.project().borrow_mut().cache.drain(..amt); 38 | } 39 | 40 | fn restore(self: Pin<&mut Self>, data: &[u8]) { 41 | self.project().borrow_mut().cache.extend_from_slice(data) 42 | } 43 | } 44 | -------------------------------------------------------------------------------- /src/io/payload/payloads/try_convert.rs: -------------------------------------------------------------------------------- 1 | use std::marker::PhantomData; 2 | use crate::io::*; 3 | 4 | #[derive(Clone)] 5 | pub struct TryConvert { 6 | prev_payload: P, 7 | action: F, 8 | _phantom: PhantomData, 9 | } 10 | 11 | impl Buildable for TryConvert where Self: PayloadAction {} 12 | impl Sendable for TryConvert where Self: PayloadAction {} 13 | impl Readable for TryConvert where Self: PayloadAction {} 14 | impl ReturnsValue for TryConvert where Self: PayloadAction {} 15 | 16 | impl TryConvert 17 | where 18 | P: ReturnsValue, 19 | F: Fn(P::ReturnType) -> Result + Clone, 20 | R: Clone 21 | { 22 | pub fn new(prev_payload: P, action: F) -> TryConvert 23 | { 24 | TryConvert { 25 | prev_payload, 26 | action, 27 | _phantom: PhantomData::default() 28 | } 29 | } 30 | } 31 | 32 | impl PayloadAction for TryConvert 33 | where 34 | P: ReturnsValue, 35 | F: Fn(P::ReturnType) -> Result + Clone, 36 | R: Clone, 37 | { 38 | type ReturnType = R; 39 | 40 | async fn execute( 41 | &self, 42 | pipe: &mut T1, 43 | ) -> Result { 44 | let result = self.prev_payload.execute(pipe).await?; 45 | (self.action)(result) 46 | } 47 | } 48 | -------------------------------------------------------------------------------- /src/io/util/timeout/async_read_timeout_ext.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::timeout::read_exact_timeout::{read_exact_timeout, ReadExactTimeout}; 2 | use crate::io::util::timeout::read_to_end_timeout::{read_to_end_timeout, ReadToEndTimeout}; 3 | use crate::io::{read_timeout, ReadTimeout}; 4 | use std::time::Duration; 5 | use tokio::io::AsyncRead; 6 | 7 | pub trait AsyncReadTimeoutExt: AsyncRead { 8 | fn read_timeout<'a>(&'a mut self, buf: &'a mut [u8], timeout: Duration) -> ReadTimeout<'a, Self> 9 | where 10 | Self: Unpin, 11 | { 12 | read_timeout(self, buf, timeout) 13 | } 14 | 15 | fn read_fill_timeout<'a>( 16 | &'a mut self, 17 | buf: &'a mut [u8], 18 | timeout: Duration, 19 | ) -> ReadExactTimeout<'a, Self> 20 | where 21 | Self: Unpin, 22 | { 23 | read_exact_timeout(self, buf, timeout, false) 24 | } 25 | 26 | fn read_exact_timeout<'a>( 27 | &'a mut self, 28 | buf: &'a mut [u8], 29 | timeout: Duration, 30 | ) -> ReadExactTimeout<'a, Self> 31 | where 32 | Self: Unpin, 33 | { 34 | read_exact_timeout(self, buf, timeout, true) 35 | } 36 | 37 | fn read_to_end_timeout<'a>( 38 | &'a mut self, 39 | buf: &'a mut Vec, 40 | timeout: Duration, 41 | throw_on_timeout: bool, 42 | ) -> ReadToEndTimeout<'a, Self> 43 | where 44 | Self: Unpin, 45 | { 46 | read_to_end_timeout(self, buf, timeout, throw_on_timeout) 47 | } 48 | } 49 | 50 | impl AsyncReadTimeoutExt for R {} 51 | -------------------------------------------------------------------------------- /src/io/payload/builder/mod.rs: -------------------------------------------------------------------------------- 1 | mod arch; 2 | mod build; 3 | mod condition; 4 | mod convert; 5 | mod dynamic_payload; 6 | mod print; 7 | mod read; 8 | mod send; 9 | 10 | use crate::io::payload::payloads; 11 | use crate::io::payload::payloads::{ 12 | Ascii, Building, Chain, Complete, DynamicPayload, Initial, ReadPayload, SendPayload, 13 | }; 14 | use crate::io::*; 15 | use crossterm::Command; 16 | use std::marker::PhantomData; 17 | 18 | #[derive(Clone)] 19 | pub struct PayloadBuilder { 20 | payload: T, 21 | _phantom_arch: PhantomData, 22 | } 23 | 24 | impl PayloadBuilder { 25 | pub fn new() -> PayloadBuilder { 26 | PayloadBuilder:: { 27 | payload: Initial {}, 28 | _phantom_arch: PhantomData::default(), 29 | } 30 | } 31 | } 32 | 33 | impl PayloadBuilder { 34 | pub fn from(payload: T) -> PayloadBuilder { 35 | PayloadBuilder:: { 36 | payload, 37 | _phantom_arch: PhantomData::default(), 38 | } 39 | } 40 | } 41 | 42 | pub trait Buildable: PayloadAction {} 43 | pub trait Sendable: PayloadAction { 44 | fn push>(self, data: T) -> impl SendCompletable 45 | where 46 | Self: Sized, 47 | { 48 | Chain::new(self, SendPayload::::new().push::(data)) 49 | } 50 | } 51 | pub trait SendCompletable: Sendable { 52 | fn complete(self) -> impl Buildable + Readable + Sendable; 53 | } 54 | 55 | pub trait Readable: PayloadAction {} 56 | pub trait ReturnsValue: PayloadAction {} 57 | -------------------------------------------------------------------------------- /src/io/payload/payloads/chain.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | 3 | #[derive(Clone)] 4 | pub struct Chain { 5 | pub payload1: P1, 6 | pub payload2: P2, 7 | } 8 | 9 | impl Buildable for Chain where Chain: PayloadAction {} 10 | impl Sendable for Chain { 11 | fn push>(self, data: T) -> impl SendCompletable 12 | where 13 | Self: Sized, 14 | { 15 | Chain::new(self.payload1, self.payload2.push::(data)) 16 | } 17 | } 18 | 19 | impl SendCompletable for Chain { 20 | fn complete(self) -> impl Buildable + Readable + Sendable { 21 | let payload2 = self.payload2.complete(); 22 | Chain::new(self.payload1, payload2) 23 | } 24 | } 25 | impl Readable for Chain where Chain: PayloadAction {} 26 | impl ReturnsValue for Chain where Chain: PayloadAction {} 27 | 28 | impl Chain { 29 | pub fn new(payload1: P1, payload2: P2) -> Chain { 30 | Chain { payload1, payload2 } 31 | } 32 | } 33 | 34 | impl PayloadAction for Chain { 35 | type ReturnType = P2::ReturnType; 36 | 37 | async fn execute( 38 | &self, 39 | pipe: &mut T, 40 | ) -> Result { 41 | let _ = self.payload1.execute(pipe).await?; 42 | self.payload2.execute(pipe).await 43 | } 44 | } 45 | -------------------------------------------------------------------------------- /src/unix/symbol.rs: -------------------------------------------------------------------------------- 1 | use crate::unix::error::ElfError; 2 | use elf::endian::AnyEndian; 3 | use elf::string_table::StringTable; 4 | use elf::ElfBytes; 5 | 6 | pub struct Symbol { 7 | pub name: String, 8 | pub shndx: u16, 9 | pub value: u64, 10 | pub size: u64, 11 | } 12 | 13 | impl Symbol { 14 | fn parse(value: &elf::symbol::Symbol, string_table: &StringTable) -> Result { 15 | let name = match value.st_name { 16 | 0 => String::new(), 17 | idx => string_table.get(idx as usize)?.to_string(), 18 | }; 19 | Ok(Symbol { 20 | name, 21 | shndx: value.st_shndx, 22 | value: value.st_value, 23 | size: value.st_size, 24 | }) 25 | } 26 | 27 | fn parse_table<'a>( 28 | table: &elf::symbol::SymbolTable<'a, AnyEndian>, 29 | string_table: &StringTable, 30 | ) -> Result, ElfError> { 31 | table 32 | .iter() 33 | .map(|s| Self::parse(&s, string_table)) 34 | .collect() 35 | } 36 | 37 | pub(crate) fn parse_symbol_table(file: &ElfBytes) -> Result, ElfError> { 38 | match file.symbol_table()? { 39 | None => Ok(Vec::new()), 40 | Some((table, string_table)) => Self::parse_table(&table, &string_table), 41 | } 42 | } 43 | 44 | pub(crate) fn parse_dynamic_symbol_table(file: &ElfBytes) -> Result, ElfError> { 45 | match file.dynamic_symbol_table()? { 46 | None => Ok(Vec::new()), 47 | Some((table, string_table)) => Self::parse_table(&table, &string_table), 48 | } 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /src/io/util/timeout/async_read_cache_timeout_ext.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::timeout::read_until_timeout::{read_until_timeout, ReadUntilTimeout}; 2 | use crate::io::{read_until_regex_timeout, AsyncCacheRead, ReadUntilRegexTimeout}; 3 | use std::time::Duration; 4 | 5 | pub trait AsyncReadCacheTimeoutExt: AsyncCacheRead { 6 | fn read_until_timeout<'a, T: AsRef<[u8]>>( 7 | &'a mut self, 8 | delim: T, 9 | buf: &'a mut Vec, 10 | timeout: Duration, 11 | ) -> ReadUntilTimeout<'a, Self, T> 12 | where 13 | Self: Unpin, 14 | { 15 | read_until_timeout(self, delim, buf, timeout) 16 | } 17 | 18 | fn read_until_regex_timeout<'a>( 19 | &'a mut self, 20 | pattern: &str, 21 | buf: &'a mut Vec, 22 | timeout: Duration, 23 | ) -> Result, regex::Error> 24 | where 25 | Self: Unpin, 26 | { 27 | read_until_regex_timeout(self, pattern, buf, timeout) 28 | } 29 | 30 | fn read_line_timeout<'a>( 31 | &'a mut self, 32 | buf: &'a mut Vec, 33 | timeout: Duration, 34 | ) -> ReadUntilTimeout<'a, Self, &'static [u8]> 35 | where 36 | Self: Unpin, 37 | { 38 | read_until_timeout(self, b"\n", buf, timeout) 39 | } 40 | 41 | fn read_line_crlf_timeout<'a>( 42 | &'a mut self, 43 | buf: &'a mut Vec, 44 | timeout: Duration, 45 | ) -> ReadUntilTimeout<'a, Self, &'static [u8]> 46 | where 47 | Self: Unpin, 48 | { 49 | read_until_timeout(self, b"\r\n", buf, timeout) 50 | } 51 | } 52 | 53 | impl AsyncReadCacheTimeoutExt for R {} 54 | -------------------------------------------------------------------------------- /src/io/payload/payloads/dynamic_payload.rs: -------------------------------------------------------------------------------- 1 | use std::marker::PhantomData; 2 | use crate::io::*; 3 | 4 | #[derive(Clone)] 5 | pub struct DynamicPayload { 6 | prev_payload: P, 7 | action: F, 8 | _phantom: PhantomData, 9 | } 10 | 11 | impl Buildable for DynamicPayload where 12 | DynamicPayload: PayloadAction 13 | { 14 | } 15 | 16 | impl Sendable for DynamicPayload where 17 | DynamicPayload: PayloadAction 18 | { 19 | } 20 | 21 | impl Readable for DynamicPayload where 22 | DynamicPayload: PayloadAction 23 | { 24 | } 25 | 26 | impl ReturnsValue for DynamicPayload where 27 | DynamicPayload: PayloadAction 28 | { 29 | } 30 | 31 | impl DynamicPayload 32 | where 33 | P: ReturnsValue, 34 | F: Fn(P::ReturnType) -> R + Clone, 35 | R: PayloadAction, 36 | { 37 | pub fn new(prev_payload: P, action: F) -> DynamicPayload 38 | { 39 | DynamicPayload { 40 | prev_payload, 41 | action, 42 | _phantom: PhantomData::default() 43 | } 44 | } 45 | } 46 | 47 | impl PayloadAction for DynamicPayload 48 | where 49 | P: PayloadAction, 50 | F: Fn(P::ReturnType) -> R + Clone, 51 | R: PayloadAction, 52 | { 53 | type ReturnType = R::ReturnType; 54 | 55 | async fn execute( 56 | &self, 57 | pipe: &mut T1, 58 | ) -> Result { 59 | let result = self.prev_payload.execute(pipe).await?; 60 | let new_payload = (self.action)(result); 61 | new_payload.execute(pipe).await 62 | } 63 | } 64 | -------------------------------------------------------------------------------- /src/io/util/timeout/read_until_timeout.rs: -------------------------------------------------------------------------------- 1 | use crate::io::ReadUntil; 2 | use crate::io::{get_deadline, read_until, AsyncCacheRead}; 3 | use pin_project_lite::pin_project; 4 | use std::future::Future; 5 | use std::io; 6 | use std::io::ErrorKind; 7 | use std::marker::PhantomPinned; 8 | use std::pin::Pin; 9 | use std::task::{Context, Poll}; 10 | use std::time::Duration; 11 | 12 | use crate::io::util::timeout::timeout; 13 | use tokio::time::Instant; 14 | 15 | pin_project! { 16 | /// The delimiter is included in the resulting vector. 17 | #[derive(Debug)] 18 | #[must_use = "futures do nothing unless you `.await` or poll them"] 19 | pub struct ReadUntilTimeout<'a, R: ?Sized, D:AsRef<[u8]>> { 20 | #[pin] 21 | read_until: ReadUntil<'a,R,D>, 22 | deadline: Instant, 23 | #[pin] 24 | _pin: PhantomPinned, 25 | } 26 | } 27 | 28 | pub(crate) fn read_until_timeout<'a, R, D: AsRef<[u8]>>( 29 | reader: &'a mut R, 30 | delimiter: D, 31 | buf: &'a mut Vec, 32 | timeout: Duration, 33 | ) -> ReadUntilTimeout<'a, R, D> 34 | where 35 | R: AsyncCacheRead + ?Sized + Unpin, 36 | { 37 | let read_until = read_until(reader, delimiter, buf); 38 | let deadline = get_deadline(timeout); 39 | ReadUntilTimeout { 40 | read_until, 41 | deadline, 42 | _pin: PhantomPinned, 43 | } 44 | } 45 | 46 | impl> Future for ReadUntilTimeout<'_, R, D> { 47 | type Output = io::Result; 48 | 49 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 50 | let me = self.project(); 51 | if *me.deadline < Instant::now() { 52 | return Err(timeout()).into(); 53 | } 54 | 55 | me.read_until.poll(cx) 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /src/io/stdio/mod.rs: -------------------------------------------------------------------------------- 1 | mod ncurses_bridge; 2 | mod shell_bridge; 3 | 4 | use crossterm::event::{KeyCode, KeyEvent, KeyEventKind, KeyModifiers}; 5 | pub use ncurses_bridge::*; 6 | pub use shell_bridge::*; 7 | use std::io; 8 | use thiserror::Error; 9 | 10 | use crate::io::pipe; 11 | use tokio::io::{AsyncRead, AsyncWrite}; 12 | use tokio::sync::mpsc; 13 | 14 | pub trait TerminalBridge { 15 | //TODO: if reader or writer does not have timeout this might get stuck 16 | async fn bridge( 17 | reader: &mut R, 18 | writer: &mut W, 19 | ); 20 | } 21 | 22 | #[derive(Error, Debug)] 23 | pub enum TerminalError { 24 | #[error("IO error: {0}")] 25 | Io(#[from] io::Error), 26 | 27 | #[error("String Send error: {0}")] 28 | StringSend(#[from] mpsc::error::SendError), 29 | 30 | #[error("Bytes Send error: {0}")] 31 | BytesSend(#[from] mpsc::error::SendError>), 32 | 33 | #[error("Recv error: {0}")] 34 | Recv(#[from] mpsc::error::TryRecvError), 35 | 36 | #[error("Pipe error: {0}")] 37 | Pipe(#[from] pipe::PipeError), 38 | 39 | #[error("format error {0}")] 40 | FmtError(#[from] std::fmt::Error), 41 | 42 | #[error("Stop signal")] 43 | Terminate, 44 | } 45 | 46 | type TerminalResult = Result; 47 | 48 | fn is_ctr_key(key_event: KeyEvent) -> bool { 49 | key_event.kind == KeyEventKind::Press && key_event.modifiers == KeyModifiers::CONTROL 50 | } 51 | 52 | pub(crate) fn is_terminate_process(key_event: KeyEvent) -> bool { 53 | return is_ctr_key(key_event) && key_event.code == KeyCode::Char('c'); 54 | } 55 | 56 | pub(crate) fn is_stop_terminal(key_event: KeyEvent) -> bool { 57 | return is_ctr_key(key_event) && key_event.code == KeyCode::Char('d'); 58 | } 59 | -------------------------------------------------------------------------------- /src/io/util/timeout/read_until_regex_timeout.rs: -------------------------------------------------------------------------------- 1 | use crate::io::{get_deadline, AsyncCacheRead}; 2 | use crate::io::{read_until_regex, ReadUntilRegex}; 3 | use pin_project_lite::pin_project; 4 | use std::future::Future; 5 | use std::io; 6 | use std::io::ErrorKind; 7 | use std::marker::PhantomPinned; 8 | use std::pin::Pin; 9 | use std::task::{Context, Poll}; 10 | use std::time::Duration; 11 | 12 | use crate::io::util::timeout::timeout; 13 | use tokio::time::Instant; 14 | 15 | pin_project! { 16 | /// The delimiter is included in the resulting vector. 17 | #[derive(Debug)] 18 | #[must_use = "futures do nothing unless you `.await` or poll them"] 19 | pub struct ReadUntilRegexTimeout<'a, R: ?Sized> { 20 | #[pin] 21 | read_until_regex: ReadUntilRegex<'a,R>, 22 | deadline: Instant, 23 | #[pin] 24 | _pin: PhantomPinned, 25 | } 26 | } 27 | 28 | pub(crate) fn read_until_regex_timeout<'a, R>( 29 | reader: &'a mut R, 30 | pattern: &str, 31 | buf: &'a mut Vec, 32 | timeout: Duration, 33 | ) -> Result, regex::Error> 34 | where 35 | R: AsyncCacheRead + ?Sized + Unpin, 36 | { 37 | let read_until_regex = read_until_regex(reader, pattern, buf)?; 38 | let deadline = get_deadline(timeout); 39 | Ok(ReadUntilRegexTimeout { 40 | read_until_regex, 41 | deadline, 42 | _pin: PhantomPinned, 43 | }) 44 | } 45 | 46 | impl Future for ReadUntilRegexTimeout<'_, R> { 47 | type Output = io::Result<(usize, usize)>; 48 | 49 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 50 | let me = self.project(); 51 | if *me.deadline < Instant::now() { 52 | return Err(timeout()).into(); 53 | } 54 | 55 | me.read_until_regex.poll(cx) 56 | } 57 | } 58 | -------------------------------------------------------------------------------- /src/io/pipe/write.rs: -------------------------------------------------------------------------------- 1 | use crate::io::{AsyncCacheRead, PipeError, PipeRead, PipeReadExt}; 2 | use crossterm::Command; 3 | use std::time::Duration; 4 | use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; 5 | 6 | pub trait PipeWrite: AsyncWrite {} 7 | 8 | impl PipeWriteExt for W {} 9 | 10 | pub trait PipeWriteExt: AsyncWrite { 11 | async fn write_line>(&mut self, text: T) -> Result 12 | where 13 | Self: Unpin, 14 | { 15 | // to_vec is used so we dont accidentally trigger 16 | // flush if user did not wrap writer into BufWriter 17 | let mut res = text.as_ref().to_vec(); 18 | res.push(b'\n'); 19 | let size = self.write(&res).await?; 20 | self.flush().await?; 21 | Ok(size) 22 | } 23 | 24 | async fn write_line_crlf>(&mut self, text: T) -> Result 25 | where 26 | Self: Unpin, 27 | { 28 | let mut res = text.as_ref().to_vec(); 29 | res.push(b'\r'); 30 | res.push(b'\n'); 31 | let size = self.write(&res).await?; 32 | self.flush().await?; 33 | Ok(size) 34 | } 35 | 36 | async fn write_flush>(&mut self, data: T) -> Result 37 | where 38 | Self: Unpin, 39 | { 40 | let size = self.write(data.as_ref()).await?; 41 | self.flush().await?; 42 | Ok(size) 43 | } 44 | 45 | async fn write_all_flush>(&mut self, data: T) -> Result<(), PipeError> 46 | where 47 | Self: Unpin, 48 | { 49 | self.write_all(data.as_ref()).await?; 50 | self.flush().await?; 51 | Ok(()) 52 | } 53 | 54 | async fn write_ansi_command(&mut self, command: T) -> Result 55 | where 56 | Self: Unpin, 57 | { 58 | let mut ansi_command = String::new(); 59 | command.write_ansi(&mut ansi_command)?; 60 | let size = self.write(ansi_command.as_bytes()).await?; 61 | self.flush().await?; 62 | Ok(size) 63 | } 64 | } 65 | -------------------------------------------------------------------------------- /src/unix/got.rs: -------------------------------------------------------------------------------- 1 | use crate::unix::error::ElfError; 2 | use crate::unix::symbol::Symbol; 3 | use elf::endian::AnyEndian; 4 | use elf::relocation::{RelIterator, RelaIterator}; 5 | use elf::{abi, ElfBytes}; 6 | use std::collections::HashMap; 7 | 8 | 9 | pub fn parse_got_from_elf( 10 | file: &ElfBytes, 11 | dynamic_symbols: &[Symbol], 12 | ) -> Result, ElfError> { 13 | let mut result = HashMap::new(); 14 | let section_table = match file.section_headers_with_strtab()? { 15 | (Some(section_table), _) => section_table, 16 | _ => { 17 | return Ok(result); 18 | } 19 | }; 20 | 21 | for section in section_table { 22 | match section.sh_type { 23 | abi::SHT_REL => { 24 | let iter = file.section_data_as_rels(§ion)?; 25 | parse_rel_iterator(iter, &dynamic_symbols, &mut result); 26 | } 27 | abi::SHT_RELA => { 28 | let iter = file.section_data_as_relas(§ion)?; 29 | parse_rela_iterator(iter, &dynamic_symbols, &mut result); 30 | } 31 | _ => {} 32 | } 33 | } 34 | 35 | Ok(result) 36 | } 37 | 38 | fn parse_rela_iterator( 39 | iter: RelaIterator, 40 | dynamic_symbols: &[Symbol], 41 | buf: &mut HashMap, 42 | ) { 43 | parse_got(iter.map(|r| (r.r_sym, r.r_offset)), dynamic_symbols, buf) 44 | } 45 | 46 | fn parse_rel_iterator( 47 | iter: RelIterator, 48 | dynamic_symbols: &[Symbol], 49 | buf: &mut HashMap, 50 | ) { 51 | parse_got(iter.map(|r| (r.r_sym, r.r_offset)), dynamic_symbols, buf) 52 | } 53 | 54 | fn parse_got( 55 | iter: impl Iterator, 56 | dynamic_symbols: &[Symbol], 57 | buf: &mut HashMap, 58 | ) { 59 | for rel in iter.filter(|r| r.0 != 0 && r.1 != 0) { 60 | match dynamic_symbols.get(rel.0 as usize) { 61 | Some(sym) if !sym.name.is_empty() => buf.insert(sym.name.clone(), rel.1), 62 | _ => continue, 63 | }; 64 | } 65 | } 66 | -------------------------------------------------------------------------------- /src/unix/elf_file.rs: -------------------------------------------------------------------------------- 1 | use crate::unix::error::ElfError; 2 | use crate::unix::symbol::Symbol; 3 | use elf::endian::AnyEndian; 4 | use elf::file::FileHeader; 5 | use elf::ElfBytes; 6 | use std::collections::HashMap; 7 | use std::ffi::OsStr; 8 | use std::path::PathBuf; 9 | use tokio::fs; 10 | use crate::unix::{parse_got_from_elf, parse_plt_from_elf}; 11 | 12 | pub struct Elf { 13 | data: Vec, 14 | header: FileHeader, 15 | symbols: HashMap, 16 | dynamic_symbols: HashMap, 17 | got: HashMap, 18 | plt: HashMap, 19 | } 20 | 21 | impl Elf { 22 | pub async fn parse>(path: &T) -> Result { 23 | let path = PathBuf::from(path); 24 | let file_data = fs::read(path).await?; 25 | let file = ElfBytes::::minimal_parse(file_data.as_slice())?; 26 | let symbols = Symbol::parse_symbol_table(&file)?; 27 | let dynamic_symbols = Symbol::parse_dynamic_symbol_table(&file)?; 28 | let got = parse_got_from_elf(&file, &dynamic_symbols)?; 29 | let plt = parse_plt_from_elf(&file, &dynamic_symbols)?; 30 | let symbols = symbols.into_iter() 31 | .map(|s| (s.name.clone(), s)) 32 | .collect(); 33 | let dynamic_symbols = dynamic_symbols.into_iter() 34 | .map(|s| (s.name.clone(), s)) 35 | .collect(); 36 | let header = file.ehdr; 37 | Ok(Elf { 38 | data: file_data, 39 | header, 40 | symbols, 41 | dynamic_symbols, 42 | got, 43 | plt, 44 | }) 45 | } 46 | 47 | pub fn plt(&self) -> &HashMap 48 | { 49 | &self.plt 50 | } 51 | 52 | pub fn got(&self) -> &HashMap 53 | { 54 | &self.got 55 | } 56 | 57 | pub fn symbols(&self) -> &HashMap 58 | { 59 | &self.symbols 60 | } 61 | 62 | pub fn dynamic_symbols(&self) -> &HashMap 63 | { 64 | &self.dynamic_symbols 65 | } 66 | 67 | pub fn header(&self) -> &FileHeader 68 | { 69 | &self.header 70 | } 71 | 72 | pub fn data(&self) -> &[u8] 73 | { 74 | &self.data 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /src/io/payload/payloads/send.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | use std::marker::PhantomData; 3 | use tokio::io::AsyncWriteExt; 4 | 5 | #[derive(Clone)] 6 | pub struct Building; 7 | #[derive(Clone)] 8 | pub struct Complete; 9 | #[derive(Clone)] 10 | pub struct SendPayload { 11 | data: Vec, 12 | _phantom: PhantomData, 13 | _phantom_arch: PhantomData, 14 | } 15 | 16 | impl Buildable for SendPayload {} 17 | impl Sendable for SendPayload {} 18 | 19 | impl Sendable for SendPayload { 20 | fn push>(self, data: T) -> impl SendCompletable 21 | where 22 | Self: Sized, 23 | { 24 | self.push_data(data) 25 | } 26 | } 27 | 28 | impl SendCompletable for SendPayload { 29 | fn complete(self) -> impl Buildable + Readable + Sendable 30 | where 31 | Self: Sized, 32 | { 33 | self.to_complete() 34 | } 35 | } 36 | impl Readable for SendPayload {} 37 | 38 | impl PayloadAction for SendPayload { 39 | type ReturnType = (); 40 | 41 | async fn execute( 42 | &self, 43 | pipe: &mut T, 44 | ) -> Result { 45 | pipe.write_all(&self.data).await?; 46 | Ok(()) 47 | } 48 | } 49 | 50 | impl PayloadAction for SendPayload { 51 | type ReturnType = (); 52 | 53 | async fn execute( 54 | &self, 55 | _pipe: &mut T, 56 | ) -> Result { 57 | unreachable!() 58 | } 59 | } 60 | 61 | impl SendPayload { 62 | pub fn new() -> Self { 63 | SendPayload { 64 | data: Vec::new(), 65 | _phantom: PhantomData::default(), 66 | _phantom_arch: PhantomData::default(), 67 | } 68 | } 69 | 70 | pub fn to_complete(self) -> SendPayload { 71 | SendPayload { 72 | data: self.data, 73 | _phantom: PhantomData::default(), 74 | _phantom_arch: PhantomData::default(), 75 | } 76 | } 77 | 78 | pub fn push_data>(mut self, data: T) -> Self { 79 | self.data.extend_from_slice(data.as_ref()); 80 | Self { 81 | data: self.data, 82 | _phantom: PhantomData::default(), 83 | _phantom_arch: PhantomData::default(), 84 | } 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ctf-pwn 2 | [![crate](https://img.shields.io/crates/v/ctf-pwn.svg)](https://crates.io/crates/ctf-pwn) 3 | 4 | Pwn utilities for Rust. 5 | 6 | ## Features 7 | 8 | ### Pipe 9 | * [x] Converting TCP Stream or Process to Pipe 10 | * [x] Conditional reading 11 | * [x] Bridging pipe to stdout/stdin 12 | * [x] Payload crafter 13 | * [x] Ansi event support for ncurses like shells 14 | 15 | ### Shell 16 | * [ ] Intel x86 17 | * [ ] Amd x64 18 | * [ ] Arm 19 | * [ ] Risc-V 20 | 21 | ### Binary Parsing 22 | * [x] Elf 23 | * [ ] PE 24 | 25 | ## Examples 26 | 27 | ### Connecting to tcp stream 28 | ```rs 29 | let mut pipe = TcpPipe::connect("127.0.0.1:1337").await?; 30 | ``` 31 | 32 | ### Spawning new process 33 | ```rs 34 | let mut pipe = ProcessPipe::from_app("ls").await?; 35 | 36 | let mut pipe = ProcessPipe::from_app_args("ls", ["-l", "-a"]).await?; 37 | ``` 38 | 39 | ### Generic loading 40 | ```rs 41 | let mut pipe = Pipe::new(stdin(), stdout()); 42 | ``` 43 | 44 | ### Reading examples 45 | ```rs 46 | let data: Vec = pipe.recv().await?; 47 | let data: Vec = pipe.recv_until("Name:", false).await?; 48 | let data: Vec = pipe.recv_until([0x01, 0x02, 0x03], false).await?; 49 | let data: String = pipe.recv_line_utf8().await?; 50 | let data: AsciiString = pipe.recv_line_ascii().await?; 51 | ``` 52 | 53 | ### Regex 54 | ```rs 55 | let data = pipe.recv_until(r"(Ok)|(Error)", true).await?; 56 | let flag = pipe.recv_regex_utf8(r"HTB\{[^\}]+\}").await?; 57 | ``` 58 | 59 | ### Interactive shell 60 | ```rs 61 | pipe.interactive_shell().await?; 62 | ``` 63 | 64 | ### Ncurses support 65 | 66 | #### Sending ansi commands 67 | ```rs 68 | pipe.write_ansi_command(ansi::Down).await?; 69 | pipe.write_ansi_command(ansi::Right).await?; 70 | pipe.write_ansi_command(ansi::Enter).await?; 71 | ``` 72 | 73 | #### Ansi event based interactive shell 74 | ```rs 75 | pipe.interactive_ansi().await?; 76 | ``` 77 | 78 | 79 | ### Payload 80 | ```rs 81 | let payload = Payload::builder() 82 | .recv_until("> ", false) 83 | .push("1") 84 | .push("\n") 85 | .send() 86 | .recv_until("Insert card's serial number: ", false) 87 | .push_line("%4919x%7$hn") 88 | .send() 89 | .recv_regex_utf8(r"HTB\{[^\}]+\}") 90 | .build(); 91 | 92 | let flag = pipe.payload(payload).await?; 93 | println!("{flag}"); 94 | ``` 95 | 96 | ### Elf 97 | ```rust 98 | let elf = Elf::parse("app_path").await?; 99 | let got: &HashMap = elf.got(); 100 | let plt: &HashMap = elf.plt(); 101 | let symbols: &HashMap = elf.symbols(); 102 | let dynamic_symbols: &HashMap = elf.dynamic_symbols(); 103 | ``` -------------------------------------------------------------------------------- /src/io/util/timeout/read_to_end_timeout.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::timeout::{get_deadline, timeout}; 2 | use pin_project_lite::pin_project; 3 | use std::future::Future; 4 | use std::io; 5 | use std::io::{BufReader, ErrorKind}; 6 | use std::marker::PhantomPinned; 7 | use std::marker::Unpin; 8 | use std::pin::Pin; 9 | use std::task::{ready, Context, Poll}; 10 | use std::time::Duration; 11 | use tokio::io::{AsyncRead, ReadBuf}; 12 | use tokio::time::Instant; 13 | 14 | pub(crate) fn read_to_end_timeout<'a, A>( 15 | reader: &'a mut A, 16 | buf: &'a mut Vec, 17 | timeout: Duration, 18 | throw_on_timeout: bool, 19 | ) -> ReadToEndTimeout<'a, A> 20 | where 21 | A: AsyncRead + Unpin + ?Sized, 22 | { 23 | let deadline = get_deadline(timeout); 24 | ReadToEndTimeout { 25 | reader, 26 | buf: buf, 27 | deadline, 28 | _pin: PhantomPinned, 29 | throw_on_timeout, 30 | } 31 | } 32 | 33 | pin_project! { 34 | #[derive(Debug)] 35 | #[must_use = "futures do nothing unless you `.await` or poll them"] 36 | pub struct ReadToEndTimeout<'a, A: ?Sized> { 37 | reader: &'a mut A, 38 | buf: &'a mut Vec, 39 | deadline: Instant, 40 | // Make this future `!Unpin` for compatibility with async trait methods. 41 | #[pin] 42 | _pin: PhantomPinned, 43 | throw_on_timeout: bool, 44 | } 45 | } 46 | 47 | impl Future for ReadToEndTimeout<'_, A> 48 | where 49 | A: AsyncRead + Unpin + ?Sized, 50 | { 51 | type Output = io::Result; 52 | 53 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 54 | let me = self.project(); 55 | let mut read_buf = [0u8; 4096]; 56 | let mut data = ReadBuf::new(&mut read_buf); 57 | loop { 58 | if *me.deadline < Instant::now() { 59 | if *me.throw_on_timeout { 60 | return Poll::Ready(Err(timeout().into())); 61 | } 62 | break; 63 | } 64 | 65 | match Pin::new(&mut *me.reader).poll_read(cx, &mut data) { 66 | Poll::Ready(Ok(_)) => {} 67 | Poll::Ready(Err(e)) if e.kind() == ErrorKind::TimedOut => { 68 | continue; 69 | } 70 | Poll::Ready(Err(e)) => { 71 | return Poll::Ready(Err(e.into())); 72 | } 73 | Poll::Pending => continue 74 | }; 75 | 76 | //eof 77 | if data.filled().len() == 0 { 78 | break; 79 | } 80 | me.buf.extend_from_slice(data.filled()); 81 | data.clear(); 82 | } 83 | 84 | Poll::Ready(Ok(me.buf.len())) 85 | } 86 | } 87 | -------------------------------------------------------------------------------- /src/io/util/timeout/read_timeout.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::timeout::{eof, get_deadline, timeout}; 2 | use pin_project_lite::pin_project; 3 | use std::future::Future; 4 | use std::io; 5 | use std::io::ErrorKind; 6 | use std::marker::PhantomPinned; 7 | use std::marker::Unpin; 8 | use std::pin::Pin; 9 | use std::task::{ready, Context, Poll}; 10 | use std::time::Duration; 11 | use tokio::io::{AsyncRead, ReadBuf}; 12 | use tokio::time::Instant; 13 | 14 | /// A future which can be used to easily read bytes until timeout or buf is fully filled 15 | pub(crate) fn read_timeout<'a, A>( 16 | reader: &'a mut A, 17 | buf: &'a mut [u8], 18 | timeout: Duration, 19 | ) -> ReadTimeout<'a, A> 20 | where 21 | A: AsyncRead + Unpin + ?Sized, 22 | { 23 | let deadline = get_deadline(timeout); 24 | ReadTimeout { 25 | reader, 26 | buf: ReadBuf::new(buf), 27 | deadline, 28 | _pin: PhantomPinned, 29 | } 30 | } 31 | 32 | pin_project! { 33 | /// Creates a future which will read exactly enough bytes to fill `buf`, 34 | /// stops if Timeout, 35 | /// returning an error if EOF is hit sooner. 36 | /// 37 | /// On success the number of bytes is returned 38 | #[derive(Debug)] 39 | #[must_use = "futures do nothing unless you `.await` or poll them"] 40 | pub struct ReadTimeout<'a, A: ?Sized> { 41 | reader: &'a mut A, 42 | buf: ReadBuf<'a>, 43 | deadline: Instant, 44 | // Make this future `!Unpin` for compatibility with async trait methods. 45 | #[pin] 46 | _pin: PhantomPinned, 47 | } 48 | } 49 | 50 | impl Future for ReadTimeout<'_, A> 51 | where 52 | A: AsyncRead + Unpin + ?Sized, 53 | { 54 | type Output = io::Result; 55 | 56 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 57 | let me = self.project(); 58 | 59 | loop { 60 | if *me.deadline < Instant::now() { 61 | return Poll::Ready(Err(timeout())); 62 | } 63 | 64 | let old_remaining = me.buf.remaining(); 65 | match Pin::new(&mut *me.reader).poll_read(cx, me.buf) { 66 | Poll::Ready(Ok(_)) => { 67 | if me.buf.remaining() == old_remaining { 68 | return Err(eof()).into(); 69 | } 70 | return Poll::Ready(Ok(me.buf.filled().len())); 71 | } 72 | Poll::Ready(Err(e)) if e.kind() == ErrorKind::TimedOut => { 73 | continue; 74 | } 75 | Poll::Ready(Err(e)) => { 76 | return Poll::Ready(Err(e.into())); 77 | } 78 | Poll::Pending => continue, 79 | }; 80 | } 81 | } 82 | } 83 | -------------------------------------------------------------------------------- /src/io/payload/builder/print.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::payload::payloads::*; 3 | use crate::io::*; 4 | use std::fmt::*; 5 | 6 | impl PayloadBuilder 7 | where 8 | T: ReturnsValue, 9 | T::ReturnType: AsRef<[u8]>, 10 | { 11 | pub fn print_bytes(self) -> PayloadBuilder, A> { 12 | PayloadBuilder::from(Print::::from(self.payload)) 13 | } 14 | } 15 | 16 | impl PayloadBuilder 17 | where 18 | T: ReturnsValue, 19 | T::ReturnType: Display, 20 | { 21 | pub fn print(self) -> PayloadBuilder, A> { 22 | PayloadBuilder::from(Print::::from(self.payload)) 23 | } 24 | } 25 | 26 | impl PayloadBuilder 27 | where 28 | T: ReturnsValue, 29 | T::ReturnType: Debug, 30 | { 31 | pub fn print_debug(self) -> PayloadBuilder, A> { 32 | PayloadBuilder::from(Print::::from(self.payload)) 33 | } 34 | } 35 | 36 | impl PayloadBuilder 37 | where 38 | T: ReturnsValue, 39 | T::ReturnType: LowerHex, 40 | { 41 | pub fn print_lower_hex(self) -> PayloadBuilder, A> { 42 | PayloadBuilder::from(Print::::from(self.payload)) 43 | } 44 | } 45 | 46 | impl PayloadBuilder 47 | where 48 | T: ReturnsValue, 49 | T::ReturnType: UpperHex, 50 | { 51 | pub fn print_upper_hex(self) -> PayloadBuilder, A> { 52 | PayloadBuilder::from(Print::::from(self.payload)) 53 | } 54 | } 55 | 56 | impl PayloadBuilder 57 | where 58 | T: ReturnsValue, 59 | T::ReturnType: Octal, 60 | { 61 | pub fn print_octal(self) -> PayloadBuilder, A> { 62 | PayloadBuilder::from(Print::::from(self.payload)) 63 | } 64 | } 65 | 66 | impl PayloadBuilder 67 | where 68 | T: ReturnsValue, 69 | T::ReturnType: Binary, 70 | { 71 | pub fn print_binary(self) -> PayloadBuilder, A> { 72 | PayloadBuilder::from(Print::::from(self.payload)) 73 | } 74 | } 75 | 76 | impl PayloadBuilder 77 | where 78 | T: ReturnsValue, 79 | T::ReturnType: Pointer, 80 | { 81 | pub fn print_pointer(self) -> PayloadBuilder, A> { 82 | PayloadBuilder::from(Print::::from(self.payload)) 83 | } 84 | } 85 | 86 | impl PayloadBuilder 87 | where 88 | T: ReturnsValue, 89 | T::ReturnType: LowerExp, 90 | { 91 | pub fn print_lower_exp(self) -> PayloadBuilder, A> { 92 | PayloadBuilder::from(Print::::from(self.payload)) 93 | } 94 | } 95 | 96 | impl PayloadBuilder 97 | where 98 | T: ReturnsValue, 99 | T::ReturnType: UpperExp, 100 | { 101 | pub fn print_upper_exp(self) -> PayloadBuilder, A> { 102 | PayloadBuilder::from(Print::::from(self.payload)) 103 | } 104 | } 105 | -------------------------------------------------------------------------------- /src/io/util/timeout/read_exact_timeout.rs: -------------------------------------------------------------------------------- 1 | use crate::io::util::timeout::{eof, get_deadline, timeout}; 2 | use pin_project_lite::pin_project; 3 | use std::future::Future; 4 | use std::io; 5 | use std::io::ErrorKind; 6 | use std::marker::PhantomPinned; 7 | use std::marker::Unpin; 8 | use std::pin::Pin; 9 | use std::task::{ready, Context, Poll}; 10 | use std::time::Duration; 11 | use tokio::io::{AsyncRead, ReadBuf}; 12 | use tokio::time::Instant; 13 | 14 | /// A future which can be used to easily read bytes until timeout or buf is fully filled 15 | pub(crate) fn read_exact_timeout<'a, A>( 16 | reader: &'a mut A, 17 | buf: &'a mut [u8], 18 | timeout: Duration, 19 | throw_on_timeout: bool, 20 | ) -> ReadExactTimeout<'a, A> 21 | where 22 | A: AsyncRead + Unpin + ?Sized, 23 | { 24 | let deadline = get_deadline(timeout); 25 | ReadExactTimeout { 26 | reader, 27 | buf: ReadBuf::new(buf), 28 | deadline, 29 | _pin: PhantomPinned, 30 | throw_on_timeout, 31 | } 32 | } 33 | 34 | pin_project! { 35 | /// Creates a future which will read exactly enough bytes to fill `buf`, 36 | /// stops if Timeout, 37 | /// returning an error if EOF is hit sooner. 38 | /// 39 | /// On success the number of bytes is returned 40 | #[derive(Debug)] 41 | #[must_use = "futures do nothing unless you `.await` or poll them"] 42 | pub struct ReadExactTimeout<'a, A: ?Sized> { 43 | reader: &'a mut A, 44 | buf: ReadBuf<'a>, 45 | deadline: Instant, 46 | // Make this future `!Unpin` for compatibility with async trait methods. 47 | #[pin] 48 | _pin: PhantomPinned, 49 | throw_on_timeout: bool, 50 | } 51 | } 52 | 53 | impl Future for ReadExactTimeout<'_, A> 54 | where 55 | A: AsyncRead + Unpin + ?Sized, 56 | { 57 | type Output = io::Result; 58 | 59 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { 60 | let me = self.project(); 61 | 62 | loop { 63 | if *me.deadline < Instant::now() { 64 | if *me.throw_on_timeout { 65 | return Poll::Ready(Err(timeout().into())); 66 | } 67 | return Poll::Ready(Ok(me.buf.filled().len())); 68 | } 69 | 70 | // if our buffer is empty, then we need to read some data to continue. 71 | let rem = me.buf.remaining(); 72 | if rem != 0 { 73 | match Pin::new(&mut *me.reader).poll_read(cx, me.buf) { 74 | Poll::Ready(Ok(_)) => {} 75 | Poll::Ready(Err(e)) if e.kind() == ErrorKind::TimedOut => { 76 | continue; 77 | } 78 | Poll::Ready(Err(e)) => { 79 | return Poll::Ready(Err(e.into())); 80 | } 81 | Poll::Pending => continue, 82 | }; 83 | if me.buf.remaining() == rem { 84 | return Err(eof()).into(); 85 | } 86 | } else { 87 | return Poll::Ready(Ok(me.buf.capacity())); 88 | } 89 | } 90 | } 91 | } 92 | -------------------------------------------------------------------------------- /src/io/util/cache/read_until_regex.rs: -------------------------------------------------------------------------------- 1 | use crate::io::AsyncCacheRead; 2 | use pin_project_lite::pin_project; 3 | use regex::bytes::*; 4 | use std::future::Future; 5 | use std::io; 6 | use std::io::ErrorKind; 7 | use std::marker::PhantomPinned; 8 | use std::pin::Pin; 9 | use std::task::{ready, Context, Poll}; 10 | use tokio::io::ReadBuf; 11 | 12 | pin_project! { 13 | /// The delimiter is included in the resulting vector. 14 | #[derive(Debug)] 15 | #[must_use = "futures do nothing unless you `.await` or poll them"] 16 | pub struct ReadUntilRegex<'a, R: ?Sized> { 17 | reader: &'a mut R, 18 | regex: Regex, 19 | buf: &'a mut Vec, 20 | internal_buf: Vec, 21 | // The number of bytes appended to buf. This can be less than buf.len() if 22 | // the buffer was not empty when the operation was started. 23 | read: usize, 24 | // Make this future `!Unpin` for compatibility with async trait methods. 25 | #[pin] 26 | _pin: PhantomPinned, 27 | } 28 | } 29 | 30 | pub(crate) fn read_until_regex<'a, R>( 31 | reader: &'a mut R, 32 | pattern: &str, 33 | buf: &'a mut Vec, 34 | ) -> Result, regex::Error> 35 | where 36 | R: AsyncCacheRead + ?Sized + Unpin, 37 | { 38 | let regex = Regex::new(pattern)?; 39 | Ok(ReadUntilRegex { 40 | reader, 41 | regex, 42 | buf, 43 | internal_buf: Vec::new(), 44 | read: 0, 45 | _pin: PhantomPinned, 46 | }) 47 | } 48 | 49 | fn eof() -> io::Error { 50 | io::Error::new(ErrorKind::UnexpectedEof, "early eof") 51 | } 52 | 53 | pub(super) fn read_until_regex_internal( 54 | mut reader: Pin<&mut R>, 55 | cx: &mut Context<'_>, 56 | regex: &mut Regex, 57 | buf: &mut Vec, 58 | internal_buf: &mut Vec, 59 | read: &mut usize, 60 | ) -> Poll> { 61 | let mut read_buf = [0u8; 4096]; 62 | let mut data = ReadBuf::new(&mut read_buf); 63 | loop { 64 | data.clear(); 65 | ready!(reader.as_mut().poll_read(cx, &mut data))?; 66 | let read_len = data.filled().len(); 67 | if read_len == 0 { 68 | return Err(eof()).into(); 69 | } 70 | *read += read_len; 71 | internal_buf.extend_from_slice(data.filled()); 72 | 73 | match regex.find(&internal_buf) { 74 | Some(m) => { 75 | let drain_index = m.end(); 76 | buf.extend_from_slice(&internal_buf[..drain_index]); 77 | let restore_data = &internal_buf[drain_index..]; 78 | reader.restore(restore_data); 79 | *read -= restore_data.len(); 80 | return Poll::Ready(Ok((buf.len(), m.len()))); 81 | } 82 | None => {} 83 | } 84 | } 85 | } 86 | 87 | impl Future for ReadUntilRegex<'_, R> { 88 | type Output = io::Result<(usize, usize)>; 89 | 90 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 91 | let me = self.project(); 92 | read_until_regex_internal( 93 | Pin::new(*me.reader), 94 | cx, 95 | me.regex, 96 | me.buf, 97 | me.internal_buf, 98 | me.read, 99 | ) 100 | } 101 | } 102 | -------------------------------------------------------------------------------- /src/io/pipe/mod.rs: -------------------------------------------------------------------------------- 1 | pub mod ansi; 2 | mod convert; 3 | mod error; 4 | mod interactive; 5 | mod read; 6 | mod write; 7 | 8 | pub use convert::*; 9 | pub use error::*; 10 | pub use interactive::*; 11 | pub use read::*; 12 | 13 | use std::io::Error; 14 | use std::pin::Pin; 15 | use std::task::{Context, Poll}; 16 | pub use write::*; 17 | 18 | use crate::io::PayloadAction; 19 | use pin_project_lite::pin_project; 20 | use std::time::Duration; 21 | use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; 22 | 23 | use super::cache::*; 24 | 25 | pin_project! { 26 | #[derive(Debug)] 27 | pub struct Pipe { 28 | #[pin] 29 | reader: CacheReader, 30 | #[pin] 31 | writer: W, 32 | timeout: Duration, 33 | block_size: usize, 34 | } 35 | } 36 | 37 | impl PipeRead for Pipe { 38 | fn get_timeout(&self) -> Duration { 39 | self.timeout 40 | } 41 | 42 | fn set_timeout(&mut self, timeout: Duration) { 43 | self.timeout = timeout; 44 | } 45 | 46 | fn get_block_size(&self) -> usize { 47 | self.block_size 48 | } 49 | 50 | fn set_block_size(&mut self, block_size: usize) { 51 | self.block_size = block_size; 52 | } 53 | } 54 | 55 | impl PipeWrite for Pipe {} 56 | 57 | impl Pipe 58 | where 59 | Self: PipeRead + PipeWrite + Send, 60 | { 61 | pub async fn payload( 62 | &mut self, 63 | payload: T, 64 | ) -> Result 65 | where 66 | Self: Unpin, 67 | { 68 | payload.execute(self).await 69 | } 70 | } 71 | 72 | impl AsyncCacheRead for Pipe { 73 | fn poll_reader( 74 | self: Pin<&mut Self>, 75 | cx: &mut Context<'_>, 76 | buf: &mut ReadBuf<'_>, 77 | ) -> Poll> { 78 | self.project().reader.poll_reader(cx, buf) 79 | } 80 | 81 | fn consume(self: Pin<&mut Self>, amt: usize) { 82 | self.project().reader.consume(amt) 83 | } 84 | 85 | fn restore(self: Pin<&mut Self>, data: &[u8]) { 86 | self.project().reader.restore(data) 87 | } 88 | } 89 | 90 | impl Pipe { 91 | const DEFAULT_TIMEOUT: Duration = Duration::from_secs(1); 92 | const DEFAULT_BLOCK_SIZE: usize = 4096; 93 | 94 | pub fn new(reader: R, writer: W) -> Pipe { 95 | Pipe { 96 | reader: CacheReader::new(reader), //CacheReader::new(timeout_reader), 97 | writer: writer, 98 | block_size: Self::DEFAULT_BLOCK_SIZE, 99 | timeout: Self::DEFAULT_TIMEOUT, 100 | } 101 | } 102 | } 103 | 104 | impl AsyncRead for Pipe { 105 | fn poll_read( 106 | self: Pin<&mut Self>, 107 | cx: &mut Context<'_>, 108 | buf: &mut ReadBuf<'_>, 109 | ) -> Poll> { 110 | let this = self.project(); 111 | this.reader.poll_read(cx, buf) 112 | } 113 | } 114 | 115 | impl AsyncWrite for Pipe { 116 | fn poll_write( 117 | self: Pin<&mut Self>, 118 | cx: &mut Context<'_>, 119 | buf: &[u8], 120 | ) -> Poll> { 121 | let this = self.project(); 122 | this.writer.poll_write(cx, buf) 123 | } 124 | 125 | fn poll_flush( 126 | self: Pin<&mut Self>, 127 | cx: &mut Context<'_>, 128 | ) -> Poll> { 129 | let this = self.project(); 130 | this.writer.poll_flush(cx) 131 | } 132 | 133 | fn poll_shutdown( 134 | self: Pin<&mut Self>, 135 | cx: &mut Context<'_>, 136 | ) -> Poll> { 137 | let this = self.project(); 138 | this.writer.poll_shutdown(cx) 139 | } 140 | } 141 | -------------------------------------------------------------------------------- /src/io/util/cache/read_until.rs: -------------------------------------------------------------------------------- 1 | use crate::io::AsyncCacheRead; 2 | use pin_project_lite::pin_project; 3 | use std::future::Future; 4 | use std::io; 5 | use std::io::ErrorKind; 6 | use std::marker::PhantomPinned; 7 | use std::pin::Pin; 8 | use std::task::{ready, Context, Poll}; 9 | use tokio::io::ReadBuf; 10 | 11 | pin_project! { 12 | /// The delimiter is included in the resulting vector. 13 | #[derive(Debug)] 14 | #[must_use = "futures do nothing unless you `.await` or poll them"] 15 | pub struct ReadUntil<'a, R: ?Sized, D:AsRef<[u8]>> { 16 | reader: &'a mut R, 17 | delimiter: D, 18 | buf: &'a mut Vec, 19 | internal_buf: Vec, 20 | // The number of bytes appended to buf. This can be less than buf.len() if 21 | // the buffer was not empty when the operation was started. 22 | read: usize, 23 | // Make this future `!Unpin` for compatibility with async trait methods. 24 | #[pin] 25 | _pin: PhantomPinned, 26 | } 27 | } 28 | 29 | pub(crate) fn read_until<'a, R, D: AsRef<[u8]>>( 30 | reader: &'a mut R, 31 | delimiter: D, 32 | buf: &'a mut Vec, 33 | ) -> ReadUntil<'a, R, D> 34 | where 35 | R: AsyncCacheRead + ?Sized + Unpin, 36 | { 37 | ReadUntil { 38 | reader, 39 | delimiter, 40 | buf, 41 | internal_buf: Vec::new(), 42 | read: 0, 43 | _pin: PhantomPinned, 44 | } 45 | } 46 | 47 | fn eof() -> io::Error { 48 | io::Error::new(ErrorKind::UnexpectedEof, "early eof") 49 | } 50 | 51 | pub(super) fn read_until_internal>( 52 | mut reader: Pin<&mut R>, 53 | cx: &mut Context<'_>, 54 | delimiter: D, 55 | buf: &mut Vec, 56 | internal_buf: &mut Vec, 57 | read: &mut usize, 58 | ) -> Poll> { 59 | let delim_len = delimiter.as_ref().len(); 60 | if delim_len == 0 { 61 | return Poll::Ready(Ok(0)); 62 | } 63 | 64 | let mut read_buf = [0u8; 4096]; 65 | let mut data = ReadBuf::new(&mut read_buf); 66 | loop { 67 | data.clear(); 68 | match ready!(reader.as_mut().poll_read(cx, &mut data)) { 69 | Ok(_) => {} 70 | Err(e) if e.kind() == ErrorKind::TimedOut => { 71 | continue; 72 | } 73 | Err(e) => { 74 | return Poll::Ready(Err(e.into())); 75 | } 76 | } 77 | let read_len = data.filled().len(); 78 | if read_len == 0 { 79 | return Err(eof()).into(); 80 | } 81 | *read += read_len; 82 | internal_buf.extend_from_slice(data.filled()); 83 | 84 | match kmp::kmp_find(delimiter.as_ref(), &internal_buf) { 85 | Some(offset) => { 86 | let drain_index = offset + delim_len; 87 | buf.extend_from_slice(&internal_buf[..drain_index]); 88 | let restore_data = &internal_buf[drain_index..]; 89 | reader.restore(restore_data); 90 | *read -= restore_data.len(); 91 | return Poll::Ready(Ok(buf.len())); 92 | } 93 | None => { 94 | if internal_buf.len() >= delim_len { 95 | let drain_range = 0..internal_buf.len() - delim_len; 96 | buf.extend_from_slice(&internal_buf[drain_range.clone()]); 97 | internal_buf.drain(drain_range); 98 | } 99 | } 100 | } 101 | } 102 | } 103 | 104 | impl> Future for ReadUntil<'_, R, D> { 105 | type Output = io::Result; 106 | 107 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 108 | let me = self.project(); 109 | read_until_internal( 110 | Pin::new(*me.reader), 111 | cx, 112 | me.delimiter, 113 | me.buf, 114 | me.internal_buf, 115 | me.read, 116 | ) 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /src/unix/plt.rs: -------------------------------------------------------------------------------- 1 | use crate::unix::error::ElfError; 2 | use crate::unix::symbol::Symbol; 3 | use elf::endian::AnyEndian; 4 | use elf::relocation::{RelIterator, RelaIterator}; 5 | use elf::{abi, ElfBytes}; 6 | use std::collections::HashMap; 7 | use elf::section::SectionHeader; 8 | 9 | 10 | fn find_plt_rel_data<'a>( 11 | file: &'a ElfBytes, 12 | ) -> Result, ElfError> { 13 | let (dynamic_table, segment_table) = match (file.dynamic()?, file.segments()) { 14 | (Some(dynamic_table), Some(segment_table)) => (dynamic_table, segment_table), 15 | _ => return Ok(None), 16 | }; 17 | 18 | let jmp_rel = match dynamic_table.iter().find(|t| t.d_tag == abi::DT_JMPREL) { 19 | Some(jmp_rel) => jmp_rel, 20 | None => return Ok(None), 21 | }; 22 | 23 | let plt_rel = match dynamic_table.iter().find(|t| t.d_tag == abi::DT_PLTREL) { 24 | Some(plt_rel) => plt_rel, 25 | None => return Ok(None), 26 | }; 27 | 28 | let plt_rel_sz = match dynamic_table.iter().find(|t| t.d_tag == abi::DT_PLTRELSZ) { 29 | Some(plt_rel_sz) => plt_rel_sz, 30 | None => return Ok(None), 31 | }; 32 | 33 | let jmp_rel_addr = jmp_rel.d_val(); 34 | 35 | let jmp_rel_header = segment_table 36 | .iter() 37 | .filter(|s| s.p_type == abi::PT_LOAD) 38 | .find(|s| jmp_rel_addr >= s.p_vaddr && (jmp_rel_addr - s.p_vaddr) < s.p_memsz); 39 | 40 | let jmp_rel_header = match jmp_rel_header { 41 | Some(jmp_rel_header) => jmp_rel_header, 42 | None => return Ok(None), 43 | }; 44 | 45 | let data = file.segment_data(&jmp_rel_header)?; 46 | let data = &data[(jmp_rel_addr - jmp_rel_header.p_vaddr) as usize..]; 47 | let data = &data[..plt_rel_sz.d_val() as usize]; 48 | 49 | let is_rela = (plt_rel.d_val() as usize) == (abi::DT_RELA as usize); 50 | 51 | Ok(Some((data, is_rela))) 52 | } 53 | 54 | pub fn parse_plt_from_elf( 55 | file: &ElfBytes, 56 | dynamic_symbols: &[Symbol], 57 | ) -> Result, ElfError> { 58 | let mut result = HashMap::new(); 59 | let (rel_data, is_rela) = match find_plt_rel_data(file)? { 60 | Some((rela_data, is_rela)) => (rela_data, is_rela), 61 | None => return Ok(result), 62 | }; 63 | 64 | let plt_section = match file.section_header_by_name(".plt")? 65 | { 66 | Some(plt_section) => plt_section, 67 | None => return Ok(result), 68 | }; 69 | 70 | match is_rela { 71 | false => parse_rel_iterator( 72 | &plt_section, 73 | RelIterator::new(file.ehdr.endianness, file.ehdr.class, rel_data), 74 | &dynamic_symbols, 75 | &mut result, 76 | ), 77 | true => parse_rela_iterator( 78 | &plt_section, 79 | RelaIterator::new(file.ehdr.endianness, file.ehdr.class, rel_data), 80 | &dynamic_symbols, 81 | &mut result, 82 | ), 83 | } 84 | 85 | Ok(result) 86 | } 87 | 88 | fn parse_rela_iterator( 89 | plt_section: &SectionHeader, 90 | iter: RelaIterator, 91 | dynamic_symbols: &[Symbol], 92 | buf: &mut HashMap, 93 | ) { 94 | parse_plt(plt_section, iter.map(|r| r.r_sym) , dynamic_symbols, buf) 95 | } 96 | 97 | fn parse_rel_iterator( 98 | plt_section: &SectionHeader, 99 | iter: RelIterator, 100 | dynamic_symbols: &[Symbol], 101 | buf: &mut HashMap, 102 | ) { 103 | parse_plt(plt_section, iter.map(|r| r.r_sym), dynamic_symbols, buf) 104 | } 105 | 106 | fn parse_plt( 107 | plt_section: &SectionHeader, 108 | iter: impl Iterator, 109 | dynamic_symbols: &[Symbol], 110 | buf: &mut HashMap, 111 | ) { 112 | for (i, r_sym) in iter.enumerate() { 113 | match dynamic_symbols.get(r_sym as usize) { 114 | Some(sym) =>{ 115 | buf.insert(sym.name.to_string(), plt_section.sh_addr + (plt_section.sh_entsize * (i + 1) as u64)) 116 | }, 117 | _ => continue, 118 | }; 119 | } 120 | } 121 | 122 | -------------------------------------------------------------------------------- /src/io/payload/payloads/read.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | use ascii::AsciiString; 3 | use std::marker::PhantomData; 4 | use std::time::Duration; 5 | 6 | #[derive(Clone)] 7 | pub struct Bytes; 8 | #[derive(Clone)] 9 | pub struct Utf8; 10 | #[derive(Clone)] 11 | pub struct Ascii; 12 | #[derive(Clone)] 13 | pub struct Interactive; 14 | 15 | impl Buildable for ReadPayload where ReadPayload: PayloadAction {} 16 | 17 | impl Readable for ReadPayload where ReadPayload: PayloadAction {} 18 | 19 | impl Sendable for ReadPayload where ReadPayload: PayloadAction {} 20 | 21 | impl ReturnsValue for ReadPayload where ReadPayload: PayloadAction {} 22 | 23 | #[derive(Clone)] 24 | pub enum ReadPayloadType { 25 | Recv(), 26 | Recvn(usize), 27 | RecvnFill(usize), 28 | RecvnExact(usize), 29 | RecvUntil(Vec, bool), 30 | RecvUntilRegex(String, bool), 31 | RecvRegex(String), 32 | RecvAll(), 33 | RecvAllTimeout(Duration, bool), 34 | RecvLine(), 35 | RecvLineCrlf(), 36 | InteractiveShell(), 37 | InteractiveAnsi(), 38 | } 39 | 40 | #[derive(Clone)] 41 | pub struct ReadPayload { 42 | read_type: ReadPayloadType, 43 | _phantom: PhantomData, 44 | } 45 | 46 | impl ReadPayload { 47 | pub fn new(read_type: ReadPayloadType) -> ReadPayload { 48 | ReadPayload { 49 | read_type, 50 | _phantom: PhantomData::default(), 51 | } 52 | } 53 | 54 | pub fn new_utf8(read_type: ReadPayloadType) -> ReadPayload { 55 | ReadPayload { 56 | read_type, 57 | _phantom: PhantomData::default(), 58 | } 59 | } 60 | 61 | pub fn new_ascii(read_type: ReadPayloadType) -> ReadPayload { 62 | ReadPayload { 63 | read_type, 64 | _phantom: PhantomData::default(), 65 | } 66 | } 67 | } 68 | 69 | async fn execute_internal( 70 | this: &ReadPayload, 71 | pipe: &mut T1, 72 | ) -> Result, PipeError> { 73 | let result = match &this.read_type { 74 | ReadPayloadType::Recv() => pipe.recv().await?, 75 | ReadPayloadType::Recvn(len) => pipe.recvn(*len).await?, 76 | ReadPayloadType::RecvnFill(len) => pipe.recvn_fill(*len).await?, 77 | ReadPayloadType::RecvnExact(len) => pipe.recvn_exact(*len).await?, 78 | ReadPayloadType::RecvUntil(delim, drop) => pipe.recv_until(delim, *drop).await?, 79 | ReadPayloadType::RecvUntilRegex(pattern, drop) => { 80 | pipe.recv_until_regex(pattern, *drop).await? 81 | } 82 | ReadPayloadType::RecvRegex(pattern) => pipe.recv_regex(pattern).await?, 83 | ReadPayloadType::RecvAll() => pipe.recv_all().await?, 84 | ReadPayloadType::RecvAllTimeout(timeout, keep_data) => { 85 | pipe.recv_all_timeout(*timeout, *keep_data).await? 86 | } 87 | ReadPayloadType::RecvLine() => pipe.recv_line().await?, 88 | ReadPayloadType::RecvLineCrlf() => pipe.recv_line_crlf().await?, 89 | ReadPayloadType::InteractiveShell() => unreachable!(), 90 | ReadPayloadType::InteractiveAnsi() => unreachable!(), 91 | }; 92 | 93 | Ok(result) 94 | } 95 | 96 | impl PayloadAction for ReadPayload { 97 | type ReturnType = Vec; 98 | 99 | async fn execute( 100 | &self, 101 | pipe: &mut T, 102 | ) -> Result { 103 | execute_internal(self, pipe).await 104 | } 105 | } 106 | 107 | impl PayloadAction for ReadPayload { 108 | type ReturnType = String; 109 | 110 | async fn execute( 111 | &self, 112 | pipe: &mut T, 113 | ) -> Result { 114 | let result = execute_internal(self, pipe).await?; 115 | Ok(String::from_utf8(result)?) 116 | } 117 | } 118 | 119 | impl PayloadAction for ReadPayload { 120 | type ReturnType = AsciiString; 121 | 122 | async fn execute( 123 | &self, 124 | pipe: &mut T, 125 | ) -> Result { 126 | let result = execute_internal(self, pipe).await?; 127 | Ok(AsciiString::from_ascii(result)?) 128 | } 129 | } 130 | 131 | impl PayloadAction for ReadPayload 132 | { 133 | type ReturnType = (); 134 | 135 | async fn execute( 136 | &self, 137 | pipe: &mut T, 138 | ) -> Result { 139 | match self.read_type { 140 | ReadPayloadType::InteractiveShell() => pipe.interactive_shell().await?, 141 | ReadPayloadType::InteractiveAnsi() => pipe.interactive_ansi().await?, 142 | _ => unreachable!() 143 | } 144 | Ok(()) 145 | } 146 | } 147 | -------------------------------------------------------------------------------- /src/io/pipe/convert/process.rs: -------------------------------------------------------------------------------- 1 | use crate::io::merge::MergedAsyncReader; 2 | use crate::io::Pipe; 3 | use std::{ffi::OsStr, process::Stdio, result}; 4 | use thiserror::*; 5 | use tokio::process::*; 6 | 7 | #[derive(Error, Debug)] 8 | pub enum ProcessPipeError { 9 | #[error("stdin not set")] 10 | StdinNotSet, 11 | #[error("stdout not set")] 12 | StdoutNotSet, 13 | #[error("stderr not set")] 14 | StdErrNotSet, 15 | #[error("io error")] 16 | IoError(#[from] tokio::io::Error), 17 | } 18 | 19 | pub type ProcessPipeResult = result::Result; 20 | 21 | pub type ProcessPipe = Pipe, ChildStdin>; 22 | pub type StdoutPipe = Pipe; 23 | pub type StderrPipe = Pipe; 24 | 25 | impl ProcessPipe { 26 | pub async fn from_app>(program: S) -> ProcessPipeResult { 27 | let command = Command::new(program); 28 | Self::spawn_command(command) 29 | } 30 | 31 | pub async fn from_app_args, I: IntoIterator>( 32 | program: S, 33 | args: I, 34 | ) -> ProcessPipeResult { 35 | let mut command = Command::new(program); 36 | let _ = command.args(args); 37 | Self::spawn_command(command) 38 | } 39 | 40 | pub fn spawn_command(mut value: Command) -> ProcessPipeResult { 41 | let process = value 42 | .stdout(Stdio::piped()) 43 | .stdin(Stdio::piped()) 44 | .stderr(Stdio::piped()) 45 | .spawn()?; 46 | 47 | process.try_into() 48 | } 49 | } 50 | 51 | impl StdoutPipe { 52 | pub async fn from_app>(program: S) -> ProcessPipeResult { 53 | let command = Command::new(program); 54 | Self::spawn_command(command) 55 | } 56 | 57 | pub async fn from_app_args, I: IntoIterator>( 58 | program: S, 59 | args: I, 60 | ) -> ProcessPipeResult { 61 | let mut command = Command::new(program); 62 | let _ = command.args(args); 63 | Self::spawn_command(command) 64 | } 65 | 66 | pub fn spawn_command(mut value: Command) -> ProcessPipeResult { 67 | let process = value.stdout(Stdio::piped()).stdin(Stdio::piped()).spawn()?; 68 | 69 | let stdin = process.stdin.ok_or(ProcessPipeError::StdinNotSet)?; 70 | let stdout = process.stdout.ok_or(ProcessPipeError::StdoutNotSet)?; 71 | Ok((stdin, stdout).into()) 72 | } 73 | } 74 | 75 | impl StderrPipe { 76 | pub async fn from_app>(program: S) -> ProcessPipeResult { 77 | let command = Command::new(program); 78 | Self::spawn_command(command) 79 | } 80 | 81 | pub async fn from_app_args, I: IntoIterator>( 82 | program: S, 83 | args: I, 84 | ) -> ProcessPipeResult { 85 | let mut command = Command::new(program); 86 | let _ = command.args(args); 87 | Self::spawn_command(command) 88 | } 89 | 90 | pub fn spawn_command(mut value: Command) -> ProcessPipeResult { 91 | let process = value.stderr(Stdio::piped()).stdin(Stdio::piped()).spawn()?; 92 | 93 | let stdin = process.stdin.ok_or(ProcessPipeError::StdinNotSet)?; 94 | let stderr = process.stderr.ok_or(ProcessPipeError::StdErrNotSet)?; 95 | Ok((stdin, stderr).into()) 96 | } 97 | } 98 | 99 | impl TryFrom for ProcessPipe { 100 | type Error = ProcessPipeError; 101 | 102 | fn try_from(value: Child) -> ProcessPipeResult { 103 | let stdin = value.stdin.ok_or(ProcessPipeError::StdinNotSet)?; 104 | let stdout = value.stdout.ok_or(ProcessPipeError::StdoutNotSet)?; 105 | let stderr = value.stderr.ok_or(ProcessPipeError::StdErrNotSet)?; 106 | Ok((stdin, stdout, stderr).into()) 107 | } 108 | } 109 | 110 | impl From<(ChildStdin, ChildStdout, ChildStderr)> for ProcessPipe { 111 | fn from(value: (ChildStdin, ChildStdout, ChildStderr)) -> Self { 112 | let (stdin, stdout, stderr) = value; 113 | let read_stream = MergedAsyncReader { 114 | stream0: stdout, 115 | stream1: stderr, 116 | }; 117 | let write_stream = stdin; 118 | Pipe::new(read_stream, write_stream) 119 | } 120 | } 121 | 122 | impl From<(ChildStdin, ChildStdout)> for StdoutPipe { 123 | fn from(value: (ChildStdin, ChildStdout)) -> Self { 124 | let (stdin, stdout) = value; 125 | let read_stream = stdout; 126 | let write_stream = stdin; 127 | Pipe::new(read_stream, write_stream) 128 | } 129 | } 130 | 131 | impl From<(ChildStdin, ChildStderr)> for StderrPipe { 132 | fn from(value: (ChildStdin, ChildStderr)) -> Self { 133 | let (stdin, stderr) = value; 134 | let read_stream = stderr; 135 | let write_stream = stdin; 136 | Pipe::new(read_stream, write_stream) 137 | } 138 | } 139 | -------------------------------------------------------------------------------- /src/io/payload/builder/send.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::*; 3 | use crossterm::Command; 4 | 5 | impl PayloadBuilder { 6 | pub fn send(self) -> PayloadBuilder { 7 | PayloadBuilder::from(self.payload.complete()) 8 | } 9 | } 10 | 11 | impl PayloadBuilder { 12 | pub fn push_ptr(self, ptr: u32) -> PayloadBuilder { 13 | self.push_u32_le(ptr) 14 | } 15 | } 16 | 17 | impl PayloadBuilder { 18 | pub fn push_ptr(self, ptr: u64) -> PayloadBuilder { 19 | self.push_u64_le(ptr) 20 | } 21 | } 22 | 23 | impl PayloadBuilder { 24 | pub fn push>(self, data: D) -> PayloadBuilder { 25 | PayloadBuilder::from(self.payload.push::(data)) 26 | } 27 | 28 | pub fn push_line>(self, data: D) -> PayloadBuilder { 29 | self.push(data).push("\n") 30 | } 31 | 32 | pub fn push_line_crlf>( 33 | self, 34 | data: D, 35 | ) -> PayloadBuilder { 36 | self.push(data).push("\r\n") 37 | } 38 | 39 | pub fn fill_byte(self, byte: u8, count: usize) -> PayloadBuilder { 40 | let mut data = Vec::new(); 41 | for _ in 0..count { 42 | data.push(byte) 43 | } 44 | self.push(data) 45 | } 46 | 47 | pub fn fill>( 48 | self, 49 | data: D, 50 | count: usize, 51 | ) -> PayloadBuilder { 52 | let mut full_data = Vec::new(); 53 | for _ in 0..count { 54 | full_data.extend_from_slice(data.as_ref()) 55 | } 56 | self.push(full_data) 57 | } 58 | 59 | pub fn push_ansi_command( 60 | self, 61 | command: D, 62 | ) -> PayloadBuilder { 63 | let mut ansi = String::new(); 64 | command.write_ansi(&mut ansi).unwrap(); 65 | self.push(ansi) 66 | } 67 | 68 | pub fn push_i8_be(self, data: i8) -> PayloadBuilder { 69 | self.push(data.to_be_bytes()) 70 | } 71 | 72 | pub fn push_i8_le(self, data: i8) -> PayloadBuilder { 73 | self.push(data.to_le_bytes()) 74 | } 75 | 76 | pub fn push_i16_be(self, data: i16) -> PayloadBuilder { 77 | self.push(data.to_be_bytes()) 78 | } 79 | 80 | pub fn push_i16_le(self, data: i16) -> PayloadBuilder { 81 | self.push(data.to_le_bytes()) 82 | } 83 | 84 | pub fn push_i32_be(self, data: i32) -> PayloadBuilder { 85 | self.push(data.to_be_bytes()) 86 | } 87 | 88 | pub fn push_i32_le(self, data: i32) -> PayloadBuilder { 89 | self.push(data.to_le_bytes()) 90 | } 91 | 92 | pub fn push_i64_be(self, data: i64) -> PayloadBuilder { 93 | self.push(data.to_be_bytes()) 94 | } 95 | 96 | pub fn push_i64_le(self, data: i64) -> PayloadBuilder { 97 | self.push(data.to_le_bytes()) 98 | } 99 | 100 | pub fn push_i128_be(self, data: i128) -> PayloadBuilder { 101 | self.push(data.to_be_bytes()) 102 | } 103 | 104 | pub fn push_i128_le(self, data: i128) -> PayloadBuilder { 105 | self.push(data.to_le_bytes()) 106 | } 107 | 108 | pub fn push_u8_be(self, data: u8) -> PayloadBuilder { 109 | self.push(data.to_be_bytes()) 110 | } 111 | 112 | pub fn push_u8_le(self, data: u8) -> PayloadBuilder { 113 | self.push(data.to_le_bytes()) 114 | } 115 | 116 | pub fn push_u16_be(self, data: u16) -> PayloadBuilder { 117 | self.push(data.to_be_bytes()) 118 | } 119 | 120 | pub fn push_u16_le(self, data: u16) -> PayloadBuilder { 121 | self.push(data.to_le_bytes()) 122 | } 123 | 124 | pub fn push_u32_be(self, data: u32) -> PayloadBuilder { 125 | self.push(data.to_be_bytes()) 126 | } 127 | 128 | pub fn push_u32_le(self, data: u32) -> PayloadBuilder { 129 | self.push(data.to_le_bytes()) 130 | } 131 | 132 | pub fn push_u64_be(self, data: u64) -> PayloadBuilder { 133 | self.push(data.to_be_bytes()) 134 | } 135 | 136 | pub fn push_u64_le(self, data: u64) -> PayloadBuilder { 137 | self.push(data.to_le_bytes()) 138 | } 139 | 140 | pub fn push_u128_be(self, data: u128) -> PayloadBuilder { 141 | self.push(data.to_be_bytes()) 142 | } 143 | 144 | pub fn push_u128_le(self, data: u128) -> PayloadBuilder { 145 | self.push(data.to_le_bytes()) 146 | } 147 | 148 | pub fn push_f32_be(self, data: f32) -> PayloadBuilder { 149 | self.push(data.to_bits().to_be_bytes()) 150 | } 151 | 152 | pub fn push_f32_le(self, data: f32) -> PayloadBuilder { 153 | self.push(data.to_bits().to_le_bytes()) 154 | } 155 | 156 | pub fn push_f64_be(self, data: f64) -> PayloadBuilder { 157 | self.push(data.to_bits().to_be_bytes()) 158 | } 159 | 160 | pub fn push_f64_le(self, data: f64) -> PayloadBuilder { 161 | self.push(data.to_bits().to_le_bytes()) 162 | } 163 | } 164 | -------------------------------------------------------------------------------- /src/io/pipe/ansi.rs: -------------------------------------------------------------------------------- 1 | use crossterm::Command; 2 | use std::fmt; 3 | 4 | #[macro_export] 5 | #[doc(hidden)] 6 | macro_rules! csi { 7 | ($( $l:expr ),*) => { concat!("\x1b", $( $l ),*) }; 8 | } 9 | 10 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 11 | pub struct Home; 12 | impl Command for Home { 13 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 14 | f.write_str(csi!("1~")) 15 | } 16 | #[cfg(windows)] 17 | fn execute_winapi(&self) -> std::io::Result<()> { 18 | Ok(()) 19 | } 20 | } 21 | 22 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 23 | pub struct Insert; 24 | impl Command for Insert { 25 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 26 | f.write_str(csi!("2~")) 27 | } 28 | #[cfg(windows)] 29 | fn execute_winapi(&self) -> std::io::Result<()> { 30 | Ok(()) 31 | } 32 | } 33 | 34 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 35 | pub struct Del; 36 | impl Command for Del { 37 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 38 | f.write_str(csi!("3~")) 39 | } 40 | #[cfg(windows)] 41 | fn execute_winapi(&self) -> std::io::Result<()> { 42 | Ok(()) 43 | } 44 | } 45 | 46 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 47 | pub struct End; 48 | impl Command for End { 49 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 50 | f.write_str(csi!("4~")) 51 | } 52 | #[cfg(windows)] 53 | fn execute_winapi(&self) -> std::io::Result<()> { 54 | Ok(()) 55 | } 56 | } 57 | 58 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 59 | pub struct PgUp; 60 | impl Command for PgUp { 61 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 62 | f.write_str(csi!("5~")) 63 | } 64 | #[cfg(windows)] 65 | fn execute_winapi(&self) -> std::io::Result<()> { 66 | Ok(()) 67 | } 68 | } 69 | 70 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 71 | pub struct PgDn; 72 | impl Command for PgDn { 73 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 74 | f.write_str(csi!("6~")) 75 | } 76 | #[cfg(windows)] 77 | fn execute_winapi(&self) -> std::io::Result<()> { 78 | Ok(()) 79 | } 80 | } 81 | 82 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 83 | pub struct F(pub u8); 84 | impl Command for F { 85 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 86 | let index = match self.0 { 87 | 0 => "10", // F0 88 | 1 => "11", // F1 89 | 2 => "12", // F2 90 | 3 => "13", // F3 91 | 4 => "14", // F4 92 | 5 => "15", // F5 93 | 6 => "17", // F6 94 | 7 => "18", // F7 95 | 8 => "19", // F8 96 | 9 => "20", // F9 97 | 10 => "21", // F10 98 | 11 => "23", // F11 99 | 12 => "24", // F12 100 | 13 => "25", // F13 101 | 14 => "26", // F14 102 | 15 => "28", // F15 103 | 16 => "29", // F16 104 | 17 => "31", // F17 105 | 18 => "32", // F18 106 | 19 => "33", // F19 107 | 20 => "34", // F20 108 | _ => panic!("Unsupported F key"), 109 | }; 110 | write!(f, csi!("{}~"), index) 111 | } 112 | #[cfg(windows)] 113 | fn execute_winapi(&self) -> std::io::Result<()> { 114 | Ok(()) 115 | } 116 | } 117 | 118 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 119 | pub struct Backspace; 120 | impl Command for Backspace { 121 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 122 | f.write_str(csi!("D"))?; 123 | f.write_str(csi!("P")) 124 | } 125 | #[cfg(windows)] 126 | fn execute_winapi(&self) -> std::io::Result<()> { 127 | Ok(()) 128 | } 129 | } 130 | 131 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 132 | pub struct Enter; 133 | impl Command for Enter { 134 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 135 | f.write_str("\n") 136 | } 137 | #[cfg(windows)] 138 | fn execute_winapi(&self) -> std::io::Result<()> { 139 | Ok(()) 140 | } 141 | } 142 | 143 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 144 | pub struct Up; 145 | impl Command for Up { 146 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 147 | f.write_str(csi!("OA")) 148 | } 149 | #[cfg(windows)] 150 | fn execute_winapi(&self) -> std::io::Result<()> { 151 | Ok(()) 152 | } 153 | } 154 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 155 | pub struct Down; 156 | impl Command for Down { 157 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 158 | f.write_str(csi!("OB")) 159 | } 160 | #[cfg(windows)] 161 | fn execute_winapi(&self) -> std::io::Result<()> { 162 | Ok(()) 163 | } 164 | } 165 | 166 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 167 | pub struct Left; 168 | impl Command for Left { 169 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 170 | f.write_str(csi!("0D")) 171 | } 172 | #[cfg(windows)] 173 | fn execute_winapi(&self) -> std::io::Result<()> { 174 | Ok(()) 175 | } 176 | } 177 | 178 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 179 | pub struct Right; 180 | impl Command for Right { 181 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 182 | f.write_str(csi!("0C")) 183 | } 184 | #[cfg(windows)] 185 | fn execute_winapi(&self) -> std::io::Result<()> { 186 | Ok(()) 187 | } 188 | } 189 | 190 | #[derive(Debug, Clone, Copy, PartialEq, Eq)] 191 | pub struct Esc; 192 | impl Command for Esc { 193 | fn write_ansi(&self, f: &mut impl fmt::Write) -> fmt::Result { 194 | f.write_str("\x1b") 195 | } 196 | #[cfg(windows)] 197 | fn execute_winapi(&self) -> std::io::Result<()> { 198 | Ok(()) 199 | } 200 | } 201 | -------------------------------------------------------------------------------- /src/io/stdio/ncurses_bridge.rs: -------------------------------------------------------------------------------- 1 | use crate::io::stdio::{is_stop_terminal, is_terminate_process, TerminalBridge, TerminalResult}; 2 | use crate::io::{ansi, TerminalError}; 3 | use crossterm::event::{KeyCode, KeyEvent, KeyEventKind}; 4 | use crossterm::*; 5 | use std::io::ErrorKind::TimedOut; 6 | use std::io::stdout; 7 | use std::sync::atomic::{AtomicBool, Ordering}; 8 | use std::sync::Arc; 9 | use std::time::Duration; 10 | use crossterm::terminal::*; 11 | use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; 12 | use tokio::join; 13 | use tokio::task::JoinError; 14 | 15 | pub struct NcursesTerminalBridge {} 16 | 17 | async fn read_task(reader: &mut R, stop_signal: Arc) -> TerminalResult<()> 18 | where 19 | R: AsyncRead + Send + Unpin, 20 | { 21 | let mut buffer = [0; 1024]; 22 | loop { 23 | if stop_signal.load(Ordering::SeqCst) { 24 | break; 25 | } 26 | 27 | let n = match reader.read(&mut buffer).await { 28 | Ok(n) if n == 0 => break, //EOF 29 | Ok(n) => n, 30 | Err(e) if e.kind() == TimedOut => continue, 31 | Err(e) => return Err(e.into()), 32 | }; 33 | 34 | let mut stdout = tokio::io::stdout(); 35 | terminal::disable_raw_mode().unwrap(); 36 | stdout.write_all(&buffer[..n]).await?; 37 | stdout.flush().await?; 38 | terminal::enable_raw_mode().unwrap(); 39 | } 40 | Ok(()) 41 | } 42 | 43 | async fn write_ansi_command( 44 | writer: &mut W, 45 | command: T, 46 | ) -> Result { 47 | let mut ansi_command = String::new(); 48 | command.write_ansi(&mut ansi_command)?; 49 | let size = writer.write(ansi_command.as_bytes()).await?; 50 | Ok(size) 51 | } 52 | 53 | async fn send_key( 54 | writer: &mut W, 55 | stop_signal: Arc, 56 | key_event: KeyEvent, 57 | ) -> TerminalResult<()> 58 | where 59 | W: AsyncWrite + Send + Unpin, 60 | { 61 | if is_terminate_process(key_event) { 62 | return Err(TerminalError::Terminate); 63 | } 64 | 65 | if is_stop_terminal(key_event) { 66 | stop_signal.store(true, Ordering::SeqCst); 67 | return Ok(()); 68 | } 69 | 70 | if key_event.kind != KeyEventKind::Press && key_event.kind != KeyEventKind::Repeat { 71 | return Ok(()); 72 | } 73 | 74 | match key_event.code { 75 | KeyCode::Char(c) => { 76 | let mut buf = [0; 4]; // UTF-8 encoding of a char may use up to 4 bytes 77 | let bytes = c.encode_utf8(&mut buf); 78 | writer.write_all(bytes.as_bytes()).await?; 79 | } 80 | KeyCode::Left => { 81 | write_ansi_command(writer, ansi::Left).await?; 82 | } 83 | KeyCode::Right => { 84 | write_ansi_command(writer, ansi::Right).await?; 85 | } 86 | KeyCode::Up => { 87 | write_ansi_command(writer, ansi::Up).await?; 88 | } 89 | KeyCode::Down => { 90 | write_ansi_command(writer, ansi::Down).await?; 91 | } 92 | KeyCode::Delete => { 93 | write_ansi_command(writer, ansi::Del).await?; 94 | } 95 | KeyCode::Backspace => { 96 | write_ansi_command(writer, ansi::Backspace).await?; 97 | } 98 | KeyCode::Esc => { 99 | write_ansi_command(writer, ansi::Esc).await?; 100 | } 101 | KeyCode::Enter => { 102 | write_ansi_command(writer, ansi::Enter).await?; 103 | } 104 | KeyCode::Home => { 105 | write_ansi_command(writer, ansi::Home).await?; 106 | } 107 | KeyCode::End => { 108 | write_ansi_command(writer, ansi::End).await?; 109 | } 110 | KeyCode::PageUp => { 111 | write_ansi_command(writer, ansi::PgUp).await?; 112 | } 113 | KeyCode::PageDown => { 114 | write_ansi_command(writer, ansi::PgDn).await?; 115 | } 116 | KeyCode::F(no) => { 117 | write_ansi_command(writer, ansi::F(no)).await?; 118 | } 119 | _ => {} 120 | }; 121 | 122 | writer.flush().await?; 123 | 124 | Ok(()) 125 | } 126 | 127 | async fn write_task(writer: &mut W, stop_signal: Arc) -> TerminalResult<()> 128 | where 129 | W: AsyncWrite + Send + Unpin, 130 | { 131 | loop { 132 | if stop_signal.load(Ordering::SeqCst) { 133 | return Ok(()); 134 | } 135 | 136 | if let Ok(true) = event::poll(Duration::from_millis(0)) { 137 | match event::read() { 138 | Ok(event::Event::Key(key_event)) => { 139 | send_key(writer, stop_signal.clone(), key_event).await?; 140 | } 141 | _ => {} 142 | } 143 | } 144 | } 145 | } 146 | 147 | impl TerminalBridge for NcursesTerminalBridge { 148 | async fn bridge( 149 | reader: &mut R, 150 | writer: &mut W, 151 | ) { 152 | let reader_ptr = reader as *mut R as usize; 153 | let writer_ptr = writer as *mut W as usize; 154 | 155 | execute!(stdout(), EnterAlternateScreen).unwrap(); 156 | terminal::enable_raw_mode().unwrap(); 157 | 158 | let stop_signal = Arc::new(AtomicBool::new(false)); 159 | let read_stop_signal = stop_signal.clone(); 160 | 161 | let reader_task = tokio::spawn(async move { 162 | let reader_ptr = reader_ptr as *mut R; 163 | let reader = unsafe { &mut *reader_ptr }; 164 | let res = read_task(reader, read_stop_signal.clone()).await; 165 | read_stop_signal.store(true, Ordering::SeqCst); 166 | res 167 | }); 168 | 169 | let writer_task = tokio::spawn(async move { 170 | let writer_ptr = writer_ptr as *mut W; 171 | let writer = unsafe { &mut *writer_ptr }; 172 | let res = write_task(writer, stop_signal.clone()).await; 173 | stop_signal.store(true, Ordering::SeqCst); 174 | res 175 | }); 176 | 177 | let (read_res, write_res) = join!(reader_task, writer_task); 178 | 179 | terminal::disable_raw_mode().unwrap(); 180 | 181 | execute!(stdout(), LeaveAlternateScreen).unwrap(); 182 | 183 | if let Ok(Err(TerminalError::Terminate)) = read_res 184 | { 185 | std::process::exit(130); //SIGINT 186 | } 187 | else if let Ok(Err(TerminalError::Terminate)) = write_res 188 | { 189 | std::process::exit(130); //SIGINT 190 | } 191 | } 192 | } 193 | -------------------------------------------------------------------------------- /src/io/payload/builder/read.rs: -------------------------------------------------------------------------------- 1 | use crate::io::payload::builder::PayloadBuilder; 2 | use crate::io::payload::payloads::{Ascii, Bytes, Chain, ReadPayload, ReadPayloadType, Utf8}; 3 | use crate::io::*; 4 | use std::time::Duration; 5 | 6 | impl PayloadBuilder { 7 | fn build_payload( 8 | self, 9 | read_type: ReadPayloadType, 10 | ) -> PayloadBuilder>, A> 11 | where 12 | ReadPayload: PayloadAction, 13 | { 14 | PayloadBuilder::from(Chain::new(self.payload, ReadPayload::new(read_type))) 15 | } 16 | 17 | pub fn recv(self) -> PayloadBuilder>, A> { 18 | self.build_payload(ReadPayloadType::Recv()) 19 | } 20 | pub fn recv_utf8(self) -> PayloadBuilder>, A> { 21 | self.build_payload(ReadPayloadType::Recv()) 22 | } 23 | pub fn recv_ascii(self) -> PayloadBuilder>, A> { 24 | self.build_payload(ReadPayloadType::Recv()) 25 | } 26 | 27 | pub fn recvn(self, len: usize) -> PayloadBuilder>, A> { 28 | self.build_payload(ReadPayloadType::Recvn(len)) 29 | } 30 | pub fn recvn_utf8(self, len: usize) -> PayloadBuilder>, A> { 31 | self.build_payload(ReadPayloadType::Recvn(len)) 32 | } 33 | pub fn recvn_ascii(self, len: usize) -> PayloadBuilder>, A> { 34 | self.build_payload(ReadPayloadType::Recvn(len)) 35 | } 36 | 37 | pub fn recvn_exact(self, len: usize) -> PayloadBuilder>, A> { 38 | self.build_payload(ReadPayloadType::RecvnExact(len)) 39 | } 40 | pub fn recvn_exact_utf8(self, len: usize) -> PayloadBuilder>, A> { 41 | self.build_payload(ReadPayloadType::RecvnExact(len)) 42 | } 43 | pub fn recvn_exact_ascii(self, len: usize) -> PayloadBuilder>, A> { 44 | self.build_payload(ReadPayloadType::RecvnExact(len)) 45 | } 46 | 47 | pub fn recvn_fill(self, len: usize) -> PayloadBuilder>, A> { 48 | self.build_payload(ReadPayloadType::RecvnFill(len)) 49 | } 50 | pub fn recvn_fill_utf8(self, len: usize) -> PayloadBuilder>, A> { 51 | self.build_payload(ReadPayloadType::RecvnFill(len)) 52 | } 53 | pub fn recvn_fill_ascii(self, len: usize) -> PayloadBuilder>, A> { 54 | self.build_payload(ReadPayloadType::RecvnFill(len)) 55 | } 56 | 57 | pub fn recv_all(self) -> PayloadBuilder>, A> { 58 | self.build_payload(ReadPayloadType::RecvAll()) 59 | } 60 | pub fn recv_all_utf8(self) -> PayloadBuilder>, A> { 61 | self.build_payload(ReadPayloadType::RecvAll()) 62 | } 63 | pub fn recv_all_ascii(self) -> PayloadBuilder>, A> { 64 | self.build_payload(ReadPayloadType::RecvAll()) 65 | } 66 | 67 | pub fn recv_all_timeout( 68 | self, 69 | timeout: Duration, 70 | keep_data: bool, 71 | ) -> PayloadBuilder>, A> { 72 | self.build_payload(ReadPayloadType::RecvAllTimeout(timeout, keep_data)) 73 | } 74 | pub fn recv_all_timeout_utf8( 75 | self, 76 | timeout: Duration, 77 | keep_data: bool, 78 | ) -> PayloadBuilder>, A> { 79 | self.build_payload(ReadPayloadType::RecvAllTimeout(timeout, keep_data)) 80 | } 81 | pub fn recv_all_timeout_ascii( 82 | self, 83 | timeout: Duration, 84 | keep_data: bool, 85 | ) -> PayloadBuilder>, A> { 86 | self.build_payload(ReadPayloadType::RecvAllTimeout(timeout, keep_data)) 87 | } 88 | 89 | pub fn recv_line(self) -> PayloadBuilder>, A> { 90 | self.build_payload(ReadPayloadType::RecvLine()) 91 | } 92 | pub fn recv_line_utf8(self) -> PayloadBuilder>, A> { 93 | self.build_payload(ReadPayloadType::RecvLine()) 94 | } 95 | pub fn recv_line_ascii(self) -> PayloadBuilder>, A> { 96 | self.build_payload(ReadPayloadType::RecvLine()) 97 | } 98 | 99 | pub fn recv_line_crlf(self) -> PayloadBuilder>, A> { 100 | self.build_payload(ReadPayloadType::RecvLineCrlf()) 101 | } 102 | pub fn recv_line_crlf_utf8(self) -> PayloadBuilder>, A> { 103 | self.build_payload(ReadPayloadType::RecvLineCrlf()) 104 | } 105 | pub fn recv_line_crlf_ascii(self) -> PayloadBuilder>, A> { 106 | self.build_payload(ReadPayloadType::RecvLineCrlf()) 107 | } 108 | 109 | pub fn recv_until>( 110 | self, 111 | data: D, 112 | drop: bool, 113 | ) -> PayloadBuilder>, A> { 114 | self.build_payload(ReadPayloadType::RecvUntil(data.as_ref().to_vec(), drop)) 115 | } 116 | pub fn recv_until_utf8>( 117 | self, 118 | data: D, 119 | drop: bool, 120 | ) -> PayloadBuilder>, A> { 121 | self.build_payload(ReadPayloadType::RecvUntil(data.as_ref().to_vec(), drop)) 122 | } 123 | pub fn recv_until_ascii>( 124 | self, 125 | data: D, 126 | drop: bool, 127 | ) -> PayloadBuilder>, A> { 128 | self.build_payload(ReadPayloadType::RecvUntil(data.as_ref().to_vec(), drop)) 129 | } 130 | 131 | pub fn recv_until_regex( 132 | self, 133 | pattern: &str, 134 | drop: bool, 135 | ) -> PayloadBuilder>, A> { 136 | self.build_payload(ReadPayloadType::RecvUntilRegex(pattern.to_string(), drop)) 137 | } 138 | pub fn recv_until_regex_utf8( 139 | self, 140 | pattern: &str, 141 | drop: bool, 142 | ) -> PayloadBuilder>, A> { 143 | self.build_payload(ReadPayloadType::RecvUntilRegex(pattern.to_string(), drop)) 144 | } 145 | pub fn recv_until_regex_ascii( 146 | self, 147 | pattern: &str, 148 | drop: bool, 149 | ) -> PayloadBuilder>, A> { 150 | self.build_payload(ReadPayloadType::RecvUntilRegex(pattern.to_string(), drop)) 151 | } 152 | 153 | pub fn recv_regex(self, pattern: &str) -> PayloadBuilder>, A> { 154 | self.build_payload(ReadPayloadType::RecvRegex(pattern.to_string())) 155 | } 156 | pub fn recv_regex_utf8(self, pattern: &str) -> PayloadBuilder>, A> { 157 | self.build_payload(ReadPayloadType::RecvRegex(pattern.to_string())) 158 | } 159 | pub fn recv_regex_ascii( 160 | self, 161 | pattern: &str, 162 | ) -> PayloadBuilder>, A> { 163 | self.build_payload(ReadPayloadType::RecvRegex(pattern.to_string())) 164 | } 165 | 166 | pub fn interactive_shell(self) -> PayloadBuilder>, A> 167 | { 168 | self.build_payload(ReadPayloadType::InteractiveShell()) 169 | } 170 | 171 | pub fn interactive_ansi(self) -> PayloadBuilder>, A> 172 | { 173 | self.build_payload(ReadPayloadType::InteractiveAnsi()) 174 | } 175 | } 176 | -------------------------------------------------------------------------------- /src/io/pipe/read.rs: -------------------------------------------------------------------------------- 1 | use std::time::Duration; 2 | 3 | use crate::io::{AsyncCacheRead, AsyncReadCacheTimeoutExt, AsyncReadTimeoutExt, PipeError}; 4 | use ascii::AsciiString; 5 | use tokio::io::{AsyncRead, AsyncReadExt}; 6 | 7 | pub trait PipeRead: AsyncRead + AsyncCacheRead { 8 | fn get_timeout(&self) -> Duration; 9 | fn set_timeout(&mut self, timeout: Duration); 10 | fn get_block_size(&self) -> usize; 11 | fn set_block_size(&mut self, block_size: usize); 12 | } 13 | 14 | impl PipeReadExt for R {} 15 | 16 | pub trait PipeReadExt: PipeRead { 17 | async fn recv(&mut self) -> Result, PipeError> 18 | where 19 | Self: Unpin, 20 | { 21 | let mut data = vec![0u8; self.get_block_size()]; 22 | let _ = self.read_timeout(&mut data, self.get_timeout()).await?; 23 | Ok(data) 24 | } 25 | 26 | async fn recv_all(&mut self) -> Result, PipeError> 27 | where 28 | Self: Unpin, 29 | { 30 | let mut data = vec![]; 31 | let _ = self.read_to_end(&mut data).await?; 32 | Ok(data) 33 | } 34 | 35 | async fn recv_all_timeout( 36 | &mut self, 37 | timeout: Duration, 38 | keep_data: bool, 39 | ) -> Result, PipeError> 40 | where 41 | Self: Unpin, 42 | { 43 | let mut data = vec![]; 44 | let _ = self 45 | .read_to_end_timeout(&mut data, timeout, !keep_data) 46 | .await?; 47 | Ok(data) 48 | } 49 | 50 | async fn recvn(&mut self, len: usize) -> Result, PipeError> 51 | where 52 | Self: Unpin, 53 | { 54 | let mut data = vec![0u8; len]; 55 | let _ = self.read_timeout(&mut data, self.get_timeout()).await?; 56 | Ok(data) 57 | } 58 | 59 | async fn recvn_fill(&mut self, len: usize) -> Result, PipeError> 60 | where 61 | Self: Unpin, 62 | { 63 | let mut data = vec![0u8; len]; 64 | let _ = self 65 | .read_fill_timeout(&mut data, self.get_timeout()) 66 | .await?; 67 | Ok(data) 68 | } 69 | 70 | async fn recvn_exact(&mut self, len: usize) -> Result, PipeError> 71 | where 72 | Self: Unpin, 73 | { 74 | let mut data = vec![0u8; len]; 75 | let _ = self 76 | .read_exact_timeout(&mut data, self.get_timeout()) 77 | .await?; 78 | Ok(data) 79 | } 80 | 81 | async fn recv_until>( 82 | &mut self, 83 | delim: T, 84 | drop: bool, 85 | ) -> Result, PipeError> 86 | where 87 | Self: Unpin, 88 | { 89 | let mut buf = Vec::new(); 90 | let delim_len = delim.as_ref().len(); 91 | self.read_until_timeout(delim, &mut buf, self.get_timeout()) 92 | .await?; 93 | if drop { 94 | buf.drain(buf.len() - delim_len..); 95 | } 96 | Ok(buf) 97 | } 98 | 99 | async fn recv_until_regex(&mut self, pattern: &str, drop: bool) -> Result, PipeError> 100 | where 101 | Self: Unpin, 102 | { 103 | let mut buf = Vec::new(); 104 | let (_, match_len) = self 105 | .read_until_regex_timeout(pattern, &mut buf, self.get_timeout())? 106 | .await?; 107 | if drop { 108 | buf.drain(buf.len() - match_len..); 109 | } 110 | Ok(buf) 111 | } 112 | 113 | async fn recv_regex(&mut self, pattern: &str) -> Result, PipeError> 114 | where 115 | Self: Unpin, 116 | { 117 | let mut buf = Vec::new(); 118 | let (_, match_len) = self 119 | .read_until_regex_timeout(pattern, &mut buf, self.get_timeout())? 120 | .await?; 121 | buf.drain(..buf.len() - match_len); 122 | Ok(buf) 123 | } 124 | 125 | async fn recv_until_regex_split( 126 | &mut self, 127 | pattern: &str, 128 | ) -> Result<(Vec, Vec), PipeError> 129 | where 130 | Self: Unpin, 131 | { 132 | let mut buf = Vec::new(); 133 | let (_, match_len) = self 134 | .read_until_regex_timeout(pattern, &mut buf, self.get_timeout())? 135 | .await?; 136 | let (data, mch) = buf.split_at(buf.len() - match_len); 137 | Ok((data.to_vec(), mch.to_vec())) 138 | } 139 | 140 | async fn recv_line(&mut self) -> Result, PipeError> 141 | where 142 | Self: Unpin, 143 | { 144 | self.recv_until(b"\n", true).await 145 | } 146 | 147 | async fn recv_line_crlf(&mut self) -> Result, PipeError> 148 | where 149 | Self: Unpin, 150 | { 151 | self.recv_until(b"\r\n", true).await 152 | } 153 | 154 | async fn recv_utf8(&mut self) -> Result 155 | where 156 | Self: Unpin, 157 | { 158 | let data = self.recv().await?; 159 | Ok(String::from_utf8(data)?) 160 | } 161 | 162 | async fn recv_until_utf8>( 163 | &mut self, 164 | delim: T, 165 | drop: bool, 166 | ) -> Result 167 | where 168 | Self: Unpin, 169 | { 170 | let data = self.recv_until(delim, drop).await?; 171 | Ok(String::from_utf8(data)?) 172 | } 173 | 174 | async fn recv_until_regex_utf8( 175 | &mut self, 176 | pattern: &str, 177 | drop: bool, 178 | ) -> Result 179 | where 180 | Self: Unpin, 181 | { 182 | let data = self.recv_until_regex(pattern, drop).await?; 183 | Ok(String::from_utf8(data)?) 184 | } 185 | 186 | async fn recv_regex_utf8(&mut self, pattern: &str) -> Result 187 | where 188 | Self: Unpin, 189 | { 190 | let data = self.recv_regex(pattern).await?; 191 | Ok(String::from_utf8(data)?) 192 | } 193 | 194 | async fn recv_line_utf8(&mut self) -> Result 195 | where 196 | Self: Unpin, 197 | { 198 | let data = self.recv_line().await?; 199 | Ok(String::from_utf8(data)?) 200 | } 201 | 202 | async fn recv_line_crlf_utf8(&mut self) -> Result 203 | where 204 | Self: Unpin, 205 | { 206 | let data = self.recv_line_crlf().await?; 207 | Ok(String::from_utf8(data)?) 208 | } 209 | 210 | async fn recv_ascii(&mut self) -> Result 211 | where 212 | Self: Unpin, 213 | { 214 | let data = self.recv().await?; 215 | Ok(AsciiString::from_ascii(data)?) 216 | } 217 | 218 | async fn recv_until_ascii>( 219 | &mut self, 220 | delim: T, 221 | drop: bool, 222 | ) -> Result 223 | where 224 | Self: Unpin, 225 | { 226 | let data = self.recv_until(delim, drop).await?; 227 | Ok(AsciiString::from_ascii(data)?) 228 | } 229 | 230 | async fn recv_until_regex_ascii( 231 | &mut self, 232 | pattern: &str, 233 | drop: bool, 234 | ) -> Result 235 | where 236 | Self: Unpin, 237 | { 238 | let data = self.recv_until_regex(pattern, drop).await?; 239 | Ok(AsciiString::from_ascii(data)?) 240 | } 241 | 242 | async fn recv_regex_ascii(&mut self, pattern: &str) -> Result 243 | where 244 | Self: Unpin, 245 | { 246 | let data = self.recv_regex(pattern).await?; 247 | Ok(AsciiString::from_ascii(data)?) 248 | } 249 | 250 | async fn recv_line_ascii(&mut self) -> Result 251 | where 252 | Self: Unpin, 253 | { 254 | let data = self.recv_line().await?; 255 | Ok(AsciiString::from_ascii(data)?) 256 | } 257 | 258 | async fn recv_line_crlf_ascii(&mut self) -> Result 259 | where 260 | Self: Unpin, 261 | { 262 | let data = self.recv_line_crlf().await?; 263 | Ok(AsciiString::from_ascii(data)?) 264 | } 265 | } 266 | -------------------------------------------------------------------------------- /src/io/payload/payloads/print.rs: -------------------------------------------------------------------------------- 1 | use crate::io::*; 2 | use std::fmt::{Binary, Debug, Display, LowerExp, LowerHex, Octal, Pointer, UpperExp, UpperHex}; 3 | use std::marker::PhantomData; 4 | 5 | #[derive(Clone)] 6 | pub struct FmtDefault; 7 | #[derive(Clone)] 8 | pub struct FmtBytes; 9 | #[derive(Clone)] 10 | pub struct FmtDebug; 11 | #[derive(Clone)] 12 | pub struct FmtLowerHex; 13 | #[derive(Clone)] 14 | pub struct FmtUpperHex; 15 | #[derive(Clone)] 16 | pub struct FmtOctal; 17 | #[derive(Clone)] 18 | pub struct FmtBinary; 19 | #[derive(Clone)] 20 | pub struct FmtPointer; 21 | #[derive(Clone)] 22 | pub struct FmtLowerExp; 23 | #[derive(Clone)] 24 | pub struct FmtUpperExp; 25 | 26 | #[derive(Clone)] 27 | pub struct Print { 28 | payload: P, 29 | _phantom: PhantomData, 30 | } 31 | 32 | impl Buildable for Print where Self: PayloadAction {} 33 | impl Sendable for Print where Self: PayloadAction {} 34 | impl Readable for Print where Self: PayloadAction {} 35 | impl ReturnsValue for Print where Self: PayloadAction {} 36 | 37 | impl

Print 38 | where 39 | P: PayloadAction, 40 | P::ReturnType: Display, 41 | { 42 | pub fn from(payload: P) -> Self { 43 | Print { 44 | payload, 45 | _phantom: PhantomData::default(), 46 | } 47 | } 48 | } 49 | 50 | impl

Print 51 | where 52 | P: PayloadAction, 53 | P::ReturnType: AsRef<[u8]>, 54 | { 55 | pub fn from(payload: P) -> Self { 56 | Print { 57 | payload, 58 | _phantom: PhantomData::default(), 59 | } 60 | } 61 | } 62 | 63 | impl

Print 64 | where 65 | P: PayloadAction, 66 | P::ReturnType: Debug, 67 | { 68 | pub fn from(payload: P) -> Self { 69 | Print { 70 | payload, 71 | _phantom: PhantomData::default(), 72 | } 73 | } 74 | } 75 | 76 | impl

Print 77 | where 78 | P: PayloadAction, 79 | P::ReturnType: LowerHex, 80 | { 81 | pub fn from(payload: P) -> Self { 82 | Print { 83 | payload, 84 | _phantom: PhantomData, 85 | } 86 | } 87 | } 88 | 89 | impl

Print 90 | where 91 | P: PayloadAction, 92 | P::ReturnType: UpperHex, 93 | { 94 | pub fn from(payload: P) -> Self { 95 | Print { 96 | payload, 97 | _phantom: PhantomData, 98 | } 99 | } 100 | } 101 | 102 | impl

Print 103 | where 104 | P: PayloadAction, 105 | P::ReturnType: Octal, 106 | { 107 | pub fn from(payload: P) -> Self { 108 | Print { 109 | payload, 110 | _phantom: PhantomData, 111 | } 112 | } 113 | } 114 | 115 | impl

Print 116 | where 117 | P: PayloadAction, 118 | P::ReturnType: Binary, 119 | { 120 | pub fn from(payload: P) -> Self { 121 | Print { 122 | payload, 123 | _phantom: PhantomData, 124 | } 125 | } 126 | } 127 | 128 | impl

Print 129 | where 130 | P: PayloadAction, 131 | P::ReturnType: Pointer, 132 | { 133 | pub fn from(payload: P) -> Self { 134 | Print { 135 | payload, 136 | _phantom: PhantomData, 137 | } 138 | } 139 | } 140 | 141 | impl

Print 142 | where 143 | P: PayloadAction, 144 | P::ReturnType: LowerExp, 145 | { 146 | pub fn from(payload: P) -> Self { 147 | Print { 148 | payload, 149 | _phantom: PhantomData, 150 | } 151 | } 152 | } 153 | 154 | impl

Print 155 | where 156 | P: PayloadAction, 157 | P::ReturnType: UpperExp, 158 | { 159 | pub fn from(payload: P) -> Self { 160 | Print { 161 | payload, 162 | _phantom: PhantomData, 163 | } 164 | } 165 | } 166 | 167 | impl

PayloadAction for Print 168 | where 169 | P: PayloadAction, 170 | P::ReturnType: Display, 171 | { 172 | type ReturnType = P::ReturnType; 173 | 174 | async fn execute( 175 | &self, 176 | pipe: &mut T, 177 | ) -> Result { 178 | let data = self.payload.execute(pipe).await?; 179 | println!("{}", data); 180 | Ok(data) 181 | } 182 | } 183 | 184 | impl

PayloadAction for Print 185 | where 186 | P: PayloadAction, 187 | P::ReturnType: AsRef<[u8]>, 188 | { 189 | type ReturnType = P::ReturnType; 190 | 191 | async fn execute( 192 | &self, 193 | pipe: &mut T, 194 | ) -> Result { 195 | let data = self.payload.execute(pipe).await?; 196 | println!("{}", hex::encode(&data)); 197 | Ok(data) 198 | } 199 | } 200 | 201 | impl

PayloadAction for Print 202 | where 203 | P: PayloadAction, 204 | P::ReturnType: Debug, 205 | { 206 | type ReturnType = P::ReturnType; 207 | 208 | async fn execute( 209 | &self, 210 | pipe: &mut T, 211 | ) -> Result { 212 | let data = self.payload.execute(pipe).await?; 213 | println!("{:?}", data); 214 | Ok(data) 215 | } 216 | } 217 | 218 | impl

PayloadAction for Print 219 | where 220 | P: PayloadAction, 221 | P::ReturnType: LowerHex, 222 | { 223 | type ReturnType = P::ReturnType; 224 | 225 | async fn execute( 226 | &self, 227 | pipe: &mut T, 228 | ) -> Result { 229 | let data = self.payload.execute(pipe).await?; 230 | println!("{:x}", data); 231 | Ok(data) 232 | } 233 | } 234 | 235 | impl

PayloadAction for Print 236 | where 237 | P: PayloadAction, 238 | P::ReturnType: UpperHex, 239 | { 240 | type ReturnType = P::ReturnType; 241 | 242 | async fn execute( 243 | &self, 244 | pipe: &mut T, 245 | ) -> Result { 246 | let data = self.payload.execute(pipe).await?; 247 | println!("{:X}", data); 248 | Ok(data) 249 | } 250 | } 251 | 252 | impl

PayloadAction for Print 253 | where 254 | P: PayloadAction, 255 | P::ReturnType: Octal, 256 | { 257 | type ReturnType = P::ReturnType; 258 | 259 | async fn execute( 260 | &self, 261 | pipe: &mut T, 262 | ) -> Result { 263 | let data = self.payload.execute(pipe).await?; 264 | println!("{:o}", data); 265 | Ok(data) 266 | } 267 | } 268 | 269 | impl

PayloadAction for Print 270 | where 271 | P: PayloadAction, 272 | P::ReturnType: Binary, 273 | { 274 | type ReturnType = P::ReturnType; 275 | 276 | async fn execute( 277 | &self, 278 | pipe: &mut T, 279 | ) -> Result { 280 | let data = self.payload.execute(pipe).await?; 281 | println!("{:b}", data); 282 | Ok(data) 283 | } 284 | } 285 | 286 | impl

PayloadAction for Print 287 | where 288 | P: PayloadAction, 289 | P::ReturnType: Pointer, 290 | { 291 | type ReturnType = P::ReturnType; 292 | 293 | async fn execute( 294 | &self, 295 | pipe: &mut T, 296 | ) -> Result { 297 | let data = self.payload.execute(pipe).await?; 298 | println!("{:p}", data); 299 | Ok(data) 300 | } 301 | } 302 | 303 | impl

PayloadAction for Print 304 | where 305 | P: PayloadAction, 306 | P::ReturnType: LowerExp, 307 | { 308 | type ReturnType = P::ReturnType; 309 | 310 | async fn execute( 311 | &self, 312 | pipe: &mut T, 313 | ) -> Result { 314 | let data = self.payload.execute(pipe).await?; 315 | println!("{:e}", data); 316 | Ok(data) 317 | } 318 | } 319 | 320 | impl

PayloadAction for Print 321 | where 322 | P: PayloadAction, 323 | P::ReturnType: UpperExp, 324 | { 325 | type ReturnType = P::ReturnType; 326 | 327 | async fn execute( 328 | &self, 329 | pipe: &mut T, 330 | ) -> Result { 331 | let data = self.payload.execute(pipe).await?; 332 | println!("{:E}", data); 333 | Ok(data) 334 | } 335 | } 336 | -------------------------------------------------------------------------------- /src/io/stdio/shell_bridge.rs: -------------------------------------------------------------------------------- 1 | use crate::io::stdio::{is_stop_terminal, is_terminate_process, TerminalBridge, TerminalResult}; 2 | 3 | use std::io::stdout; 4 | use std::io::ErrorKind::TimedOut; 5 | use std::io::Write; 6 | use std::sync::atomic::{AtomicBool, Ordering}; 7 | use std::sync::Arc; 8 | use std::time::Duration; 9 | use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; 10 | 11 | use crossterm::cursor::{DisableBlinking, EnableBlinking, MoveTo}; 12 | use crossterm::event::{ 13 | DisableBracketedPaste, EnableBracketedPaste, KeyCode, KeyEvent, KeyEventKind, 14 | }; 15 | use crossterm::style::Print; 16 | use crossterm::terminal::{Clear, ClearType, EnterAlternateScreen, LeaveAlternateScreen}; 17 | use crossterm::*; 18 | use tokio::join; 19 | use tokio::sync::mpsc::error::TryRecvError; 20 | use tokio::sync::mpsc::{channel, Receiver, Sender}; 21 | use crate::io::{AsyncReadTimeoutExt, TerminalError}; 22 | 23 | pub struct ShellTerminalBridge {} 24 | 25 | struct StdoutState<'a, W: AsyncWrite + Unpin> { 26 | text: String, 27 | start_position: (u16, u16), 28 | cursor_position: (u16, u16), 29 | //TODO: Handle resize correctly 30 | current_dimensions: (u16, u16), 31 | writer: &'a mut W, 32 | stop_signal: Arc, 33 | } 34 | 35 | impl<'a, W: AsyncWrite + Unpin> StdoutState<'a, W> { 36 | pub fn new(writer: &mut W, stop_signal: Arc) -> TerminalResult> { 37 | let cursor_position = cursor::position()?; 38 | let current_dimensions = terminal::size()?; 39 | Ok(StdoutState { 40 | text: String::new(), 41 | start_position: cursor_position, 42 | cursor_position, 43 | current_dimensions, 44 | writer, 45 | stop_signal, 46 | }) 47 | } 48 | 49 | pub async fn insert(&mut self, key_event: KeyEvent) -> TerminalResult<()> { 50 | if is_terminate_process(key_event) { 51 | return Err(TerminalError::Terminate); 52 | } 53 | 54 | if is_stop_terminal(key_event) { 55 | self.stop_signal.store(true, Ordering::SeqCst); 56 | return Ok(()); 57 | } 58 | 59 | if key_event.kind != KeyEventKind::Press && key_event.kind != KeyEventKind::Repeat { 60 | return Ok(()); 61 | } 62 | 63 | match key_event.code { 64 | KeyCode::Char(c) => self.insert_char(c)?, 65 | KeyCode::Left => self.decrement_cursor()?, 66 | KeyCode::Right => self.increment_cursor()?, 67 | KeyCode::Backspace => self.backspace()?, 68 | KeyCode::Delete => self.del()?, 69 | KeyCode::Enter => self.send_data().await?, 70 | KeyCode::Home => self.home()?, 71 | KeyCode::End => self.end()?, 72 | _ => {} 73 | }; 74 | 75 | Ok(()) 76 | } 77 | 78 | fn get_cursor_relative_index(&self) -> usize { 79 | let (start_x, start_y) = self.start_position; 80 | let (end_x, end_y) = self.cursor_position; 81 | let (w, _h) = self.current_dimensions; 82 | 83 | let full_lines = if end_y > start_y { 84 | end_y - start_y - 1 85 | } else { 86 | 0 87 | }; 88 | let last_line_chars = if end_y > start_y { 89 | end_x 90 | } else { 91 | end_x - start_x 92 | }; 93 | 94 | full_lines as usize * w as usize + last_line_chars as usize 95 | } 96 | 97 | fn set_cursor_relative_index(&mut self, index: usize) -> TerminalResult<()> { 98 | if index > self.text.len() { 99 | return Ok(()); 100 | } 101 | 102 | let (start_x, start_y) = self.start_position; 103 | let (w, _) = self.current_dimensions; 104 | 105 | let lines_down = index / w as usize; 106 | let new_y = start_y + lines_down as u16; 107 | 108 | let new_x = (start_x as usize + index % w as usize) as u16; 109 | 110 | let (final_x, final_y) = if new_x >= w { 111 | (new_x - w, new_y + 1) 112 | } else { 113 | (new_x, new_y) 114 | }; 115 | 116 | self.set_cursor_position(final_x, final_y) 117 | } 118 | 119 | fn increment_cursor(&mut self) -> TerminalResult<()> { 120 | let index = self.get_cursor_relative_index(); 121 | if index >= self.text.len() { 122 | return Ok(()); 123 | } 124 | self.set_cursor_relative_index(index + 1) 125 | } 126 | 127 | fn decrement_cursor(&mut self) -> TerminalResult<()> { 128 | let index = self.get_cursor_relative_index(); 129 | if index <= 0 { 130 | return Ok(()); 131 | } 132 | self.set_cursor_relative_index(index - 1) 133 | } 134 | 135 | pub fn set_cursor_position(&mut self, x: u16, y: u16) -> TerminalResult<()> { 136 | self.cursor_position = (x, y); 137 | execute!(stdout(), MoveTo(x, y))?; 138 | Ok(()) 139 | } 140 | 141 | pub fn print(&mut self, data: &[u8]) -> TerminalResult<()> { 142 | self.clear()?; 143 | terminal::disable_raw_mode().unwrap(); 144 | stdout().write_all(data)?; 145 | stdout().flush()?; 146 | terminal::enable_raw_mode().unwrap(); 147 | self.start_position = cursor::position()?; 148 | self.cursor_position = self.start_position; 149 | self.redraw()?; 150 | Ok(()) 151 | } 152 | 153 | pub fn insert_str(&mut self, text: &str) -> TerminalResult<()> { 154 | let index = self.get_cursor_relative_index(); 155 | self.text.insert_str(index, text); 156 | self.redraw()?; 157 | self.set_cursor_relative_index(index + text.len())?; 158 | Ok(()) 159 | } 160 | 161 | pub fn insert_char(&mut self, c: char) -> TerminalResult<()> { 162 | let index = self.get_cursor_relative_index(); 163 | self.text.insert(index, c); 164 | self.redraw()?; 165 | self.set_cursor_relative_index(index + 1)?; 166 | Ok(()) 167 | } 168 | 169 | fn del(&mut self) -> TerminalResult<()> { 170 | let index = self.get_cursor_relative_index(); 171 | if index >= self.text.len() { 172 | return Ok(()); 173 | } 174 | let _ = self.text.remove(index); 175 | self.redraw()?; 176 | self.set_cursor_relative_index(index)?; 177 | Ok(()) 178 | } 179 | 180 | fn backspace(&mut self) -> TerminalResult<()> { 181 | let index = self.get_cursor_relative_index(); 182 | if index <= 0 { 183 | return Ok(()); 184 | } 185 | let _ = self.text.remove(index - 1); 186 | self.redraw()?; 187 | self.set_cursor_relative_index(index - 1)?; 188 | Ok(()) 189 | } 190 | 191 | pub fn home(&mut self) -> TerminalResult<()> { 192 | self.set_cursor_relative_index(0) 193 | } 194 | 195 | pub fn end(&mut self) -> TerminalResult<()> { 196 | self.set_cursor_relative_index(self.text.len()) 197 | } 198 | 199 | pub async fn send_data(&mut self) -> TerminalResult<()> { 200 | self.end()?; 201 | self.redraw()?; 202 | let (_, y) = cursor::position()?; 203 | println!(); 204 | self.set_cursor_position(0, y + 1)?; 205 | self.start_position = self.cursor_position; 206 | let mut text = self.text.clone(); 207 | self.text.clear(); 208 | text.push('\n'); 209 | 210 | self.writer.write_all(text.as_bytes()).await?; 211 | self.writer.flush().await?; 212 | Ok(()) 213 | } 214 | 215 | fn clear(&mut self) -> TerminalResult<()> { 216 | let (start_x, start_y) = self.start_position; 217 | self.set_cursor_position(start_x, start_y)?; 218 | 219 | execute!( 220 | stdout(), 221 | MoveTo(start_x, start_y), 222 | Clear(ClearType::FromCursorDown) 223 | )?; 224 | Ok(()) 225 | } 226 | 227 | pub fn redraw(&mut self) -> TerminalResult<()> { 228 | self.clear()?; 229 | let (start_x, start_y) = self.start_position; 230 | 231 | execute!(stdout(), MoveTo(start_x, start_y), Print(&self.text))?; 232 | Ok(()) 233 | } 234 | } 235 | 236 | async fn read_task( 237 | reader: &mut R, 238 | stop_signal: Arc, 239 | sender: Sender>, 240 | ) -> TerminalResult<()> 241 | where 242 | R: AsyncRead + Send + Unpin, 243 | { 244 | let mut buffer = [0; 1024]; 245 | loop { 246 | if stop_signal.load(Ordering::SeqCst) { 247 | break; 248 | } 249 | 250 | let n = match reader.read_timeout(&mut buffer, Duration::from_secs(1)).await { 251 | Ok(n) if n == 0 => break, //EOF 252 | Ok(n) => n, 253 | Err(e) if e.kind() == TimedOut =>{ 254 | continue; 255 | } 256 | Err(e) => return Err(e.into()), 257 | }; 258 | 259 | sender.send(buffer[..n].to_vec()).await? 260 | } 261 | Ok(()) 262 | } 263 | 264 | async fn write_task( 265 | writer: &mut W, 266 | stop_signal: Arc, 267 | mut receiver: Receiver>, 268 | ) -> TerminalResult<()> 269 | where 270 | W: AsyncWrite + Send + Unpin, 271 | { 272 | let mut stdout = StdoutState::new(writer, stop_signal.clone())?; 273 | 274 | loop { 275 | if stop_signal.load(Ordering::SeqCst) { 276 | return Ok(()); 277 | } 278 | 279 | match receiver.try_recv() { 280 | Ok(data) => { 281 | stdout.print(&data)?; 282 | } 283 | Err(TryRecvError::Empty) => {} 284 | Err(TryRecvError::Disconnected) => return Err(TryRecvError::Disconnected.into()), 285 | } 286 | 287 | if let Ok(true) = event::poll(Duration::from_millis(0)) { 288 | match event::read() { 289 | Ok(event::Event::Key(key_event)) => { 290 | stdout.insert(key_event).await?; 291 | } 292 | Ok(event::Event::Resize(_width, _height)) => { 293 | //TODO: recalculate cursor after resize 294 | } 295 | Ok(event::Event::Paste(text)) => { 296 | stdout.insert_str(&text)?; 297 | } 298 | _ => {} 299 | } 300 | } 301 | } 302 | } 303 | 304 | impl TerminalBridge for ShellTerminalBridge { 305 | async fn bridge( 306 | reader: &mut R, 307 | writer: &mut W, 308 | ) { 309 | let reader_ptr = reader as *mut R as usize; 310 | let writer_ptr = writer as *mut W as usize; 311 | 312 | let (rx, tx) = channel(100); 313 | 314 | let _ = execute!(stdout(), EnterAlternateScreen); 315 | terminal::enable_raw_mode().unwrap(); 316 | let _ = execute!(stdout(), EnableBlinking, EnableBracketedPaste); 317 | 318 | let stop_signal = Arc::new(AtomicBool::new(false)); 319 | let read_stop_signal = stop_signal.clone(); 320 | 321 | let reader_task = tokio::spawn(async move { 322 | let reader_ptr = reader_ptr as *mut R; 323 | let reader = unsafe { &mut *reader_ptr }; 324 | let res = read_task(reader, read_stop_signal.clone(), rx).await; 325 | read_stop_signal.store(true, Ordering::SeqCst); 326 | res 327 | }); 328 | 329 | let writer_task = tokio::spawn(async move { 330 | let writer_ptr = writer_ptr as *mut W; 331 | let writer = unsafe { &mut *writer_ptr }; 332 | let res = write_task(writer, stop_signal.clone(), tx).await; 333 | stop_signal.store(true, Ordering::SeqCst); 334 | res 335 | }); 336 | 337 | let (read_res, write_res) = join!(reader_task, writer_task); 338 | 339 | let _ = execute!(stdout(), DisableBlinking, DisableBracketedPaste); 340 | 341 | terminal::disable_raw_mode().unwrap(); 342 | 343 | let _ = execute!(stdout(), LeaveAlternateScreen); 344 | 345 | if let Ok(Err(TerminalError::Terminate)) = read_res 346 | { 347 | std::process::exit(130); //SIGINT 348 | } 349 | else if let Ok(Err(TerminalError::Terminate)) = write_res 350 | { 351 | std::process::exit(130); //SIGINT 352 | } 353 | } 354 | } 355 | --------------------------------------------------------------------------------