From a1ff624f2a7eb9e0cdf81f1ea9eecac90fef26ee Mon Sep 17 00:00:00 2001 From: C0ffeeCode Date: Sat, 28 Sep 2024 20:48:50 +0200 Subject: [PATCH] Refactor: Split `pty` and `shell` out of `channel` --- src/channel.rs | 129 ++++++++----------------------------------- src/channel/pty.rs | 76 +++++++++++++++++++++++++ src/channel/shell.rs | 49 ++++++++++++++++ src/connection.rs | 10 ++-- 4 files changed, 153 insertions(+), 111 deletions(-) create mode 100644 src/channel/pty.rs create mode 100644 src/channel/shell.rs diff --git a/src/channel.rs b/src/channel.rs index 0bd09f8..e254971 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -1,12 +1,14 @@ -use std::fs::{File, OpenOptions}; -use std::io::{self, Read, Write}; -use std::os::unix::io::{FromRawFd, IntoRawFd, RawFd}; -use std::os::unix::process::CommandExt; -use std::path::PathBuf; -use std::process::{self, Stdio}; -use std::thread::{self, JoinHandle}; +mod pty; +mod shell; -use crate::sys; +use std::fs::File; +use std::io::{self, Write}; +use std::os::unix::io::RawFd; +use std::path::PathBuf; +use std::process; +use std::thread::JoinHandle; + +pub use pty::PtyConfig; pub type ChannelId = u32; @@ -25,23 +27,18 @@ pub struct Channel { #[derive(Debug)] pub enum ChannelRequest { - Pty { - term: String, - chars: u16, - rows: u16, - pixel_width: u16, - pixel_height: u16, - modes: Vec, - }, + Pty(PtyConfig), Shell, } impl Channel { pub fn new( - id: ChannelId, peer_id: ChannelId, peer_window_size: u32, - max_packet_size: u32 - ) -> Channel { - Channel { + id: ChannelId, + peer_id: ChannelId, + peer_window_size: u32, + max_packet_size: u32, + ) -> Self { + Self { id, peer_id, process: None, @@ -66,99 +63,19 @@ impl Channel { self.max_packet_size } - pub fn request(&mut self, request: ChannelRequest) { - match request - { - ChannelRequest::Pty { - chars, - rows, - pixel_width, - pixel_height, - .. - } => { - let (master_fd, tty_path) = sys::getpty(); - - sys::set_winsize( - master_fd, - chars, - rows, - pixel_width, - pixel_height, - ); - - self.read_thread = Some(thread::spawn(move || { - #[cfg(target_os = "redox")] - let master2 = unsafe { syscall::dup(master_fd as usize, &[]).unwrap_or(!0) }; - #[cfg(not(target_os = "redox"))] - let master2 = unsafe { libc::dup(master_fd) }; - - println!("dup result: {}", master2 as u32); - let mut master = unsafe { File::from_raw_fd(master2 as i32) }; - loop { - use std::str::from_utf8_unchecked; - let mut buf = [0; 4096]; - let count = master.read(&mut buf).unwrap(); - // This is weird. - // An error is thrown&unwrapped here (panic) - // but yet it continues to function properly - if count == 0 { - break; - } - println!("Read: {}", unsafe { - from_utf8_unchecked(&buf[0..count]) - }); - } - - println!("Quitting read thread."); - })); - - self.pty = Some((master_fd, tty_path)); - self.master = Some(unsafe { File::from_raw_fd(master_fd) }); - } - ChannelRequest::Shell => { - 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 handle_request(&mut self, request: ChannelRequest) { + match request { + ChannelRequest::Pty(ref pty) => self.setup_tty(pty), + ChannelRequest::Shell => self.setup_shell(), } debug!("Channel Request: {:?}", request); } - pub fn data(&mut self, data: &[u8]) -> io::Result<()> { + 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 { + } else { Ok(()) } } diff --git a/src/channel/pty.rs b/src/channel/pty.rs new file mode 100644 index 0000000..21016b0 --- /dev/null +++ b/src/channel/pty.rs @@ -0,0 +1,76 @@ +use std::{fs::File, io::Read, os::fd::FromRawFd, thread}; + +use crate::sys; + +use super::Channel; + +#[derive(Debug)] +pub struct PtyConfig { + pub term: String, + pub chars: u16, + pub rows: u16, + pub pixel_width: u16, + pub pixel_height: u16, + pub modes: Vec, +} + +impl Channel { + /// Allocates and sets up new PTY. + /// TODO: consider what to do if there already is a PTY. + /// + /// Also does incorrect error handling internally + /// on I/O errors and occurs if a process exits. + pub fn setup_tty( + &mut self, + PtyConfig { + term, + chars, + rows, + pixel_width, + pixel_height, + modes, + }: &PtyConfig, + ) { + let (master_fd, tty_path) = sys::getpty(); + + sys::set_winsize(master_fd, *chars, *rows, *pixel_width, *pixel_height); + + self.read_thread = Some(thread::spawn(move || { + #[cfg(target_os = "redox")] + let master2 = + unsafe { syscall::dup(master_fd as usize, &[]).unwrap_or(!0) }; + #[cfg(not(target_os = "redox"))] + let master2 = unsafe { libc::dup(master_fd) }; + + println!("dup result: {}", master2 as u32); + let mut master = unsafe { File::from_raw_fd(master2 as i32) }; + loop { + use std::str::from_utf8_unchecked; + + let mut buf = [0; 4096]; + let count = match master.read(&mut buf) { + Ok(o) => o, + Err(e) => { + warn!("Error occured, ignoring: {}", e); + 1 // TODO + }, + }; + + // This is weird. + // An error is thrown&unwrapped here (panic) + // but yet it continues to function properly + if count == 0 { + break; + } + println!("Read: {}", unsafe { + from_utf8_unchecked(&buf[0..count]) + }); + } + + println!("Quitting read thread."); + })); + + self.pty = Some((master_fd, tty_path)); + self.master = Some(unsafe { File::from_raw_fd(master_fd) }); + } +} diff --git a/src/channel/shell.rs b/src/channel/shell.rs new file mode 100644 index 0000000..41a9c4b --- /dev/null +++ b/src/channel/shell.rs @@ -0,0 +1,49 @@ +use std::{ + fs::OpenOptions, + os::{ + fd::{FromRawFd, IntoRawFd}, + unix::process::CommandExt, + }, + process::{self, Stdio}, +}; + +use crate::sys; + +use super::Channel; + +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(); + } + } +} diff --git a/src/connection.rs b/src/connection.rs index 9128099..5edddb1 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use rand::distributions::Standard; -use crate::channel::{Channel, ChannelId, ChannelRequest}; +use crate::channel::{Channel, ChannelId, ChannelRequest, PtyConfig}; use crate::encryption::{AesCtr, Decryptor, Encryption}; use crate::error::{ConnectionError, ConnectionResult as Result}; use crate::key_exchange::{KexResult, KeyExchange}; @@ -315,21 +315,21 @@ impl Connection { let want_reply = reader.read_bool()?; let request = match &*name { - "pty-req" => Some(ChannelRequest::Pty { + "pty-req" => Some(ChannelRequest::Pty(PtyConfig { term: reader.read_utf8()?, chars: reader.read_uint32()? as u16, rows: reader.read_uint32()? as u16, pixel_width: reader.read_uint32()? as u16, pixel_height: reader.read_uint32()? as u16, modes: reader.read_string()?, - }), + })), "shell" => Some(ChannelRequest::Shell), _ => None, }; if let Some(request) = request { let channel = self.channels.get_mut(&channel_id).unwrap(); - channel.request(request); + channel.handle_request(request); } else { warn!("Unkown channel request {}", name); } @@ -349,7 +349,7 @@ impl Connection { let data = reader.read_string()?; let channel = self.channels.get_mut(&channel_id).unwrap(); - channel.data(data.as_slice())?; + channel.write_data(data.as_slice())?; Ok(None) }