From d9af15ae5a2ddd89dfab1e6a960e5c1d42170223 Mon Sep 17 00:00:00 2001 From: C0ffeeCode Date: Sat, 5 Oct 2024 16:28:25 +0200 Subject: [PATCH] Fix (shell): Send `std(out|err)` and read to `stdin` non-blocking TTY remains not functional --- .vscode/settings.json | 2 +- src/channel.rs | 38 +++++++++++++++- src/channel/pty.rs | 7 ++- src/channel/shell.rs | 101 ++++++++++++++++++++++++++---------------- src/connection.rs | 50 +++++++++++++++++---- src/sys/unix.rs | 12 ++++- 6 files changed, 157 insertions(+), 53 deletions(-) diff --git a/.vscode/settings.json b/.vscode/settings.json index 9eed10f..6cc59c4 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -2,7 +2,7 @@ "lldb.displayFormat": "auto", "lldb.showDisassembly": "never", // "rust-analyzer.cargo.allTargets": true, - "rust-analyzer.cargo.target": "x86_64-unknown-redox", + // "rust-analyzer.cargo.target": "x86_64-unknown-redox", // "rust-analyzer.check.targets": [ // "x86_64-unknown-redox", // "x86_64-unknown-linux-gnu" diff --git a/src/channel.rs b/src/channel.rs index e254971..fda158a 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -2,22 +2,25 @@ mod pty; mod shell; use std::fs::File; -use std::io::{self, Write}; +use std::io::{self, BufRead, Read, Write}; use std::os::unix::io::RawFd; use std::path::PathBuf; use std::process; use std::thread::JoinHandle; pub use pty::PtyConfig; +use shell::PipeContainer; pub type ChannelId = u32; #[derive(Debug)] +/// TODO: Split into channel types, i.e. `PtyChannel` and `ShellChannel` pub struct Channel { id: ChannelId, peer_id: ChannelId, process: Option, pty: Option<(RawFd, PathBuf)>, + pipes: Option, master: Option, window_size: u32, peer_window_size: u32, @@ -44,6 +47,7 @@ impl Channel { process: None, master: None, pty: None, + pipes: None, window_size: peer_window_size, peer_window_size, max_packet_size, @@ -71,12 +75,44 @@ impl Channel { debug!("Channel Request: {:?}", request); } + /// Writes `data` **to** this channel; pub fn write_data(&mut self, data: &[u8]) -> io::Result<()> { if let Some(ref mut master) = self.master { master.write_all(data)?; master.flush() + } else if let Some(PipeContainer { ref mut stdin, .. }) = self.pipes { + stdin.write_all(data)?; + stdin.flush() } else { Ok(()) } } + + /// Reads data **from** this channel and writes it **to** the buffer/. + pub fn read_pty_master(&mut self, buf: &mut [u8]) -> io::Result { + if let Some(ref mut master) = self.master { + master.read(buf) + } else { + Ok(0) + } + } + + /// Reads data **from** this channel and writes it **to** the buffer/. + pub fn read_stdout(&mut self, buf: &mut String) -> io::Result { + if let Some(pipes) = &mut self.pipes { + let res_len = pipes.stdout.read_line(buf); + res_len + } else { + Ok(0) + } + } + + pub fn read_stderr(&mut self, buf: &mut String) -> io::Result { + if let Some(pipes) = &mut self.pipes { + let res_len = pipes.stderr.read_line(buf); + res_len + } else { + Ok(0) + } + } } diff --git a/src/channel/pty.rs b/src/channel/pty.rs index 21016b0..f45833d 100644 --- a/src/channel/pty.rs +++ b/src/channel/pty.rs @@ -1,4 +1,4 @@ -use std::{fs::File, io::Read, os::fd::FromRawFd, thread}; +use std::{fs::File, io::{Read, Write}, os::fd::FromRawFd, thread}; use crate::sys; @@ -62,9 +62,8 @@ impl Channel { if count == 0 { break; } - println!("Read: {}", unsafe { - from_utf8_unchecked(&buf[0..count]) - }); + let data_read = unsafe { from_utf8_unchecked(&buf[0..count]) }; + println!("Read: {}", data_read); } println!("Quitting read thread."); diff --git a/src/channel/shell.rs b/src/channel/shell.rs index 41a9c4b..c7c2a1d 100644 --- a/src/channel/shell.rs +++ b/src/channel/shell.rs @@ -1,49 +1,76 @@ use std::{ - fs::OpenOptions, - os::{ - fd::{FromRawFd, IntoRawFd}, + fs::OpenOptions, io::{stdin, BufReader}, os::{ + fd::{FromRawFd, IntoRawFd, RawFd}, unix::process::CommandExt, - }, - process::{self, Stdio}, + }, path::PathBuf, process::{self, ChildStderr, ChildStdin, ChildStdout, Stdio} }; use crate::sys; use super::Channel; +#[derive(Debug)] +pub struct PipeContainer { + pub stdin: ChildStdin, + pub stdout: BufReader, + pub stderr: BufReader, +} + impl Channel { - pub fn setup_shell(&self) { - if let Some((_, tty_path)) = self.pty.as_ref() { - let stdin = OpenOptions::new() - .read(true) - .write(true) - .open(tty_path) - .unwrap() - .into_raw_fd(); - - let stdout = OpenOptions::new() - .read(true) - .write(true) - .open(tty_path) - .unwrap() - .into_raw_fd(); - - let stderr = OpenOptions::new() - .read(true) - .write(true) - .open(tty_path) - .unwrap() - .into_raw_fd(); - - unsafe { - process::Command::new("login") - .stdin(Stdio::from_raw_fd(stdin)) - .stdout(Stdio::from_raw_fd(stdout)) - .stderr(Stdio::from_raw_fd(stderr)) - .pre_exec(sys::before_exec) - } - .spawn() - .unwrap(); + pub fn setup_shell(&mut self) { + match self.pty.as_ref() { + Some((_, tty_path)) => with_tty(tty_path), + None => { + let pipes = without_tty(); + #[cfg(unix)] + use crate::sys::non_blockify_reader; + non_blockify_reader(pipes.stdout.get_ref()); + non_blockify_reader(pipes.stderr.get_ref()); + self.pipes = Some(pipes); + }, } } } + +fn without_tty() -> PipeContainer { + let proc = unsafe { + process::Command::new("/bin/sh") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .pre_exec(sys::before_exec) + .spawn() + .unwrap() + }; + + PipeContainer { + stdin: proc.stdin.unwrap(), + stdout: BufReader::with_capacity(1, proc.stdout.unwrap()), + stderr: BufReader::with_capacity(1, proc.stderr.unwrap()), + } +} + +fn with_tty(tty_path: &PathBuf) { + let stdin = open_tty(tty_path, true, true); + let stdout = open_tty(tty_path, true, true); + let stderr = open_tty(tty_path, true, true); + + unsafe { + process::Command::new("/bin/sh") + .stdin(Stdio::from_raw_fd(stdin)) + .stdout(Stdio::from_raw_fd(stdout)) + .stderr(Stdio::from_raw_fd(stderr)) + .pre_exec(sys::before_exec) + } + .spawn() + .unwrap(); +} + +fn open_tty(tty_path: &PathBuf, read: bool, write: bool) -> RawFd { + OpenOptions::new() + .read(read) + .write(write) + .open(tty_path) + .unwrap() + .into_raw_fd() +} diff --git a/src/connection.rs b/src/connection.rs index 5edddb1..ef3700e 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -62,16 +62,23 @@ impl Connection { } } - pub fn run(&mut self, stream: &mut S) -> Result<()> { + pub fn run(&mut self, stream: &mut S) -> Result<()> { self.send_id(stream)?; self.read_id(stream)?; + crate::sys::non_blockify_reader(stream); let mut reader = BufReader::new(stream); loop { - let packet = self.recv(&mut reader)?; - let response = self.process(packet)?; + std::thread::sleep(std::time::Duration::from_millis(1)); // TODO: REMOVE: debugging + let response = match self.recv(&mut reader) { + Ok(packet) => self.process_packet(packet)?, + Err(ConnectionError::Io(ref io_err)) if io_err.kind() == io::ErrorKind::WouldBlock => None, + Err(err) => return Err(err), + }; + // let response = self.process_packet(packet)?; + // Take back ownership of the stream let mut stream = reader.get_mut(); if let Some(packet) = response { @@ -80,6 +87,32 @@ impl Connection { // Send additional packets from the queue let mut packets: Vec = self.tx_queue.drain(..).collect(); + // Include packets from command output + // TODO: Consider pushing to `tx_queue`? + for (_, channel) in self.channels.iter_mut() { + let mut buf = String::with_capacity(1024); + if channel.read_pty_master(unsafe { buf.as_mut_vec() }).unwrap_or(0) != 0 { + let mut packet = Packet::new(MessageType::ChannelData); + packet.write_uint32(channel.id())?; + packet.write_string(&buf)?; + info!("AHAA! {:?}", &packet); + packets.push(packet) + } + if channel.read_stdout(&mut buf).unwrap_or(0) != 0 { + let mut packet = Packet::new(MessageType::ChannelData); + packet.write_uint32(channel.id())?; + packet.write_string(&buf)?; + packets.push(packet) + } + if channel.read_stderr(&mut buf).unwrap_or(0) != 0 { + let mut packet = Packet::new(MessageType::ChannelExtendedData); + packet.write_uint32(channel.id())?; + packet.write_uint32(1)?; // Data Type Code for stderr + packet.write_string(&buf)?; + packets.push(packet) + } + } + // Actually send the queued packets for packet in packets.drain(..) { self.send(&mut stream, packet)?; } @@ -89,11 +122,10 @@ impl Connection { fn recv(&mut self, mut stream: &mut dyn Read) -> Result { let packet = if let Some((ref mut c2s, _)) = self.encryption { let mut decryptor = Decryptor::new(&mut **c2s, &mut stream); - Packet::read_from(&mut decryptor)? - } - else { - Packet::read_from(&mut stream)? - }; + Packet::read_from(&mut decryptor) + } else { + Packet::read_from(&mut stream) + }?; if let Some((ref mut mac, _)) = self.mac { let mut sig = vec![0; mac.size()]; @@ -197,7 +229,7 @@ impl Connection { Ok(key) } - pub fn process(&mut self, packet: Packet) -> Result> { + pub fn process_packet(&mut self, packet: Packet) -> Result> { match packet.msg_type() { MessageType::KexInit => self.kex_init(packet), MessageType::NewKeys => self.new_keys(packet), diff --git a/src/sys/unix.rs b/src/sys/unix.rs index f5d7797..f3874f8 100644 --- a/src/sys/unix.rs +++ b/src/sys/unix.rs @@ -1,5 +1,5 @@ use std::io::Result; -use std::os::unix::io::RawFd; +use std::os::fd::{AsRawFd, RawFd}; use std::path::PathBuf; pub fn before_exec() -> Result<()> { @@ -69,3 +69,13 @@ pub fn getpty() -> (RawFd, PathBuf) { }; (master_fd, tty_path) } + +/// Sets the file descriptor to non-blocking mode. +/// This uses the **`unsafe`** keyword +pub fn non_blockify_reader(obj: &impl AsRawFd) { + use libc::{fcntl, F_GETFL, F_SETFL, O_NONBLOCK}; + + let fd = obj.as_raw_fd(); + let flags = unsafe { fcntl(fd, F_GETFL) }; + unsafe { fcntl(fd, F_SETFL, flags | O_NONBLOCK) }; +}