From 1e652077744d298c6692b425f75c980c177cebb3 Mon Sep 17 00:00:00 2001 From: Jordan Doyle Date: Sat, 17 Feb 2024 14:17:07 +0000 Subject: [PATCH] Implement cat command --- pisshoff-server/src/command.rs | 14 +++++++++----- pisshoff-server/src/command/cat.rs | 242 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ pisshoff-server/src/command/ls.rs | 57 ++++++++++++++++++++++++++++++++++++++++++++------------- pisshoff-server/src/file_system.rs | 156 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------ 4 files changed, 445 insertions(+), 24 deletions(-) create mode 100644 pisshoff-server/src/command/cat.rs diff --git a/pisshoff-server/src/command.rs b/pisshoff-server/src/command.rs index 2ef082b..8d33613 100644 --- a/pisshoff-server/src/command.rs +++ b/pisshoff-server/src/command.rs @@ -1,3 +1,4 @@ +mod cat; mod echo; mod exit; mod ls; @@ -6,13 +7,14 @@ mod scp; mod uname; mod whoami; -use crate::server::{ConnectionState, ThrusshSession}; +use std::{borrow::Cow, fmt::Debug}; + use async_trait::async_trait; use itertools::Either; -use std::borrow::Cow; -use std::fmt::Debug; use thrussh::ChannelId; +use crate::server::{ConnectionState, ThrusshSession}; + #[derive(Debug)] pub enum CommandResult { /// Wait for stdin @@ -146,7 +148,8 @@ define_commands! { Pwd(pwd::Pwd) = b"pwd", Scp(scp::Scp) = b"scp", Uname(uname::Uname) = b"uname", - Whoami(whoami::Whoami) = b"whoami" + Whoami(whoami::Whoami) = b"whoami", + Cat(cat::Cat) = b"cat" } #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -170,9 +173,10 @@ fn argparse(args: &[String]) -> impl Iterator> { #[cfg(test)] mod test { - use super::Arg; use test_case::test_case; + use super::Arg; + #[test_case("-a", &[Arg::Short('a')]; "single short parameter")] #[test_case("-abc", &[Arg::Short('a'), Arg::Short('b'), Arg::Short('c')]; "multiple short parameter")] #[test_case("-a --long operand -b -", &[Arg::Short('a'), Arg::Long("long"), Arg::Operand("operand"), Arg::Short('b'), Arg::Operand("-")]; "full hit")] diff --git a/pisshoff-server/src/command/cat.rs b/pisshoff-server/src/command/cat.rs new file mode 100644 index 0000000..efd3d0b --- /dev/null +++ b/pisshoff-server/src/command/cat.rs @@ -0,0 +1,242 @@ +use std::{collections::VecDeque, path::Path}; + +use async_trait::async_trait; +use thrussh::ChannelId; + +use crate::{ + command::{Command, CommandResult}, + server::{ConnectionState, ThrusshSession}, +}; + +#[derive(Debug, Clone)] +pub struct Cat { + remaining_params: VecDeque, + status: u32, +} + +impl Cat { + fn run( + mut self, + connection: &mut ConnectionState, + channel: ChannelId, + session: &mut S, + ) -> CommandResult { + while let Some(param) = self.remaining_params.pop_front() { + if param == "-" { + return CommandResult::ReadStdin(self); + } + + match connection.file_system().read(Path::new(¶m)) { + Ok(content) => { + session.data(channel, content.to_vec().into()); + } + Err(e) => { + self.status = 1; + // TODO: stderr + eprintln!("{e}"); + session.data(channel, format!("cat: {param}: {e}").into()); + } + } + } + + CommandResult::Exit(self.status) + } +} + +#[async_trait] +impl Command for Cat { + async fn new( + connection: &mut ConnectionState, + params: &[String], + channel: ChannelId, + session: &mut S, + ) -> CommandResult { + let this = Self { + remaining_params: params.to_vec().into(), + status: 0, + }; + + if params.is_empty() { + CommandResult::ReadStdin(this) + } else { + this.run(connection, channel, session) + } + } + + async fn stdin( + self, + connection: &mut ConnectionState, + channel: ChannelId, + data: &[u8], + session: &mut S, + ) -> CommandResult { + session.data(channel, data.to_vec().into()); + self.run(connection, channel, session) + } +} + +#[cfg(test)] +mod test { + use std::path::Path; + + use mockall::predicate::{self, always}; + + use crate::{ + command::{cat::Cat, Command, CommandResult}, + server::{ + test::{fake_channel_id, predicate::eq_string}, + ConnectionState, MockThrusshSession, + }, + }; + + #[tokio::test] + async fn no_args() { + let mut session = MockThrusshSession::default(); + + let out = Cat::new( + &mut ConnectionState::mock(), + [].as_slice(), + fake_channel_id(), + &mut session, + ) + .await; + + assert!(matches!(out, CommandResult::ReadStdin(_)), "{out:?}"); + } + + #[tokio::test] + async fn file_args_with_missing() { + let mut session = MockThrusshSession::default(); + let mut state = ConnectionState::mock(); + + state.file_system().mkdirall(Path::new("/rootdir")).unwrap(); + + state + .file_system() + .write(Path::new("a"), "hello".as_bytes().into()) + .unwrap(); + state + .file_system() + .write(Path::new("/rootdir/c"), "world".as_bytes().into()) + .unwrap(); + + session + .expect_data() + .once() + .with(always(), eq_string("hello")) + .returning(|_, _| ()); + + session + .expect_data() + .once() + .with(always(), eq_string("cat: b: No such file or directory")) + .returning(|_, _| ()); + + session + .expect_data() + .once() + .with(always(), eq_string("world")) + .returning(|_, _| ()); + + let out = Cat::new( + &mut state, + ["a".to_string(), "b".to_string(), "/rootdir/c".to_string()].as_slice(), + fake_channel_id(), + &mut session, + ) + .await; + + assert!(matches!(out, CommandResult::Exit(1)), "{out:?}"); + } + + #[tokio::test] + async fn file_args() { + let mut session = MockThrusshSession::default(); + let mut state = ConnectionState::mock(); + + state + .file_system() + .write(Path::new("a"), "hello".as_bytes().into()) + .unwrap(); + state + .file_system() + .write(Path::new("b"), "world".as_bytes().into()) + .unwrap(); + + session + .expect_data() + .once() + .with(always(), eq_string("hello")) + .returning(|_, _| ()); + + session + .expect_data() + .once() + .with(always(), eq_string("world")) + .returning(|_, _| ()); + + let out = Cat::new( + &mut state, + ["a".to_string(), "b".to_string()].as_slice(), + fake_channel_id(), + &mut session, + ) + .await; + + assert!(matches!(out, CommandResult::Exit(0)), "{out:?}"); + } + + #[tokio::test] + async fn stdin() { + let mut session = MockThrusshSession::default(); + let mut state = ConnectionState::mock(); + + state + .file_system() + .write(Path::new("a"), "hello".as_bytes().into()) + .unwrap(); + + state + .file_system() + .write(Path::new("b"), "world".as_bytes().into()) + .unwrap(); + + session + .expect_data() + .once() + .with(always(), eq_string("hello")) + .returning(|_, _| ()); + + session + .expect_data() + .once() + .with(always(), eq_string("the whole")) + .returning(|_, _| ()); + + session + .expect_data() + .once() + .with(always(), eq_string("world")) + .returning(|_, _| ()); + + let out = Cat::new( + &mut state, + ["a".to_string(), "-".to_string(), "b".to_string()].as_slice(), + fake_channel_id(), + &mut session, + ) + .await + .unwrap_stdin(); + + let out = out + .stdin( + &mut state, + fake_channel_id(), + "the whole".as_bytes(), + &mut session, + ) + .await; + + assert!(matches!(out, CommandResult::Exit(0)), "{out:?}"); + } +} diff --git a/pisshoff-server/src/command/ls.rs b/pisshoff-server/src/command/ls.rs index 81278bd..75e4e87 100644 --- a/pisshoff-server/src/command/ls.rs +++ b/pisshoff-server/src/command/ls.rs @@ -1,10 +1,12 @@ +use std::{fmt::Write, path::Path}; + +use async_trait::async_trait; +use thrussh::ChannelId; + use crate::{ command::{Command, CommandResult}, server::{ConnectionState, ThrusshSession}, }; -use async_trait::async_trait; -use std::fmt::Write; -use thrussh::ChannelId; #[derive(Debug, Clone)] pub struct Ls {} @@ -17,33 +19,55 @@ impl Command for Ls { channel: ChannelId, session: &mut S, ) -> CommandResult { + let mut error = false; + let resp = if params.is_empty() { - connection.file_system().ls(None).join(" ") + match connection.file_system().ls(None) { + Ok(v) => v.join(" "), + Err(e) => { + error = true; + format!("ls: {}: {e}", connection.file_system().pwd().display()) + } + } } else if params.len() == 1 { - connection + match connection .file_system() - .ls(Some(params.first().unwrap())) - .join(" ") + .ls(Some(Path::new(params.first().unwrap()))) + { + Ok(v) => v.join(" "), + Err(e) => { + error = true; + format!("ls: {}: {e}", params.first().unwrap()) + } + } } else { let mut out = String::new(); for dir in params { if !out.is_empty() { - out.push_str("\n\n"); + out.push('\n'); } - write!(out, "{dir}:").unwrap(); - out.push_str(&connection.file_system().ls(Some(dir)).join(" ")); + match connection.file_system().ls(Some(Path::new(dir))) { + Ok(v) => { + write!(out, "{dir}:\n{}", v.join(" ")).unwrap(); + } + Err(e) => { + error = true; + write!(out, "ls: {dir}: {e}").unwrap(); + } + } } out }; if !resp.is_empty() { + let resp = resp.trim(); session.data(channel, format!("{resp}\n").into()); } - CommandResult::Exit(0) + CommandResult::Exit(u32::from(error)) } async fn stdin( @@ -59,6 +83,10 @@ impl Command for Ls { #[cfg(test)] mod test { + use std::path::Path; + + use mockall::predicate::always; + use crate::{ command::{ls::Ls, Command, CommandResult}, server::{ @@ -66,7 +94,6 @@ mod test { ConnectionState, MockThrusshSession, }, }; - use mockall::predicate::always; #[tokio::test] async fn empty_pwd() { @@ -93,8 +120,12 @@ mod test { .with(always(), eq_string("a:\n\nb:\n")) .returning(|_, _| ()); + let mut state = ConnectionState::mock(); + state.file_system().mkdirall(Path::new("/root/a")).unwrap(); + state.file_system().mkdirall(Path::new("/root/b")).unwrap(); + let out = Ls::new( - &mut ConnectionState::mock(), + &mut state, ["a".to_string(), "b".to_string()].as_slice(), fake_channel_id(), &mut session, diff --git a/pisshoff-server/src/file_system.rs b/pisshoff-server/src/file_system.rs index 5f96385..3882137 100644 --- a/pisshoff-server/src/file_system.rs +++ b/pisshoff-server/src/file_system.rs @@ -1,25 +1,57 @@ #![allow(dead_code)] -use std::path::{Path, PathBuf}; +use std::{ + borrow::Cow, + collections::{btree_map::Entry, BTreeMap}, + fmt::{Display, Formatter}, + path::{Path, PathBuf}, +}; /// A fake file system, stored in memory only active for the current session. pub struct FileSystem { pwd: PathBuf, home: PathBuf, + data: Tree, +} + +pub enum Tree { + Directory(BTreeMap>), + File(Box<[u8]>), } impl FileSystem { pub fn new(user: &str) -> Self { let pwd = if user == "root" { - PathBuf::new().join("/root") + PathBuf::from("/root") } else { - PathBuf::new().join("/home").join(user) + PathBuf::from("/home").join(user) }; - Self { + let mut this = Self { home: pwd.clone(), pwd, + data: Tree::Directory(BTreeMap::new()), + }; + + let _res = this.mkdirall(&this.pwd.clone()); + this + } + + pub fn mkdirall(&mut self, path: &Path) -> Result<(), LsError> { + let mut tree = &mut self.data; + + for c in path { + match tree { + Tree::Directory(d) => { + tree = d + .entry(c.to_str().unwrap().to_string()) + .or_insert_with(|| Box::new(Tree::Directory(BTreeMap::new()))); + } + Tree::File(_) => return Err(LsError::FileExists), + } } + + Ok(()) } pub fn cd(&mut self, v: Option<&str>) { @@ -34,8 +66,120 @@ impl FileSystem { &self.pwd } + pub fn read(&self, path: &Path) -> Result<&[u8], LsError> { + let canonical = self.pwd().join(path); + let mut tree = &self.data; + + for c in &canonical { + match tree { + Tree::Directory(d) => { + tree = d + .get(c.to_str().unwrap()) + .ok_or(LsError::NoSuchFileOrDirectory)?; + } + Tree::File(_) => { + return Err(LsError::NotDirectory); + } + } + } + + match tree { + Tree::Directory(_) => Err(LsError::IsADirectory), + Tree::File(content) => Ok(content), + } + } + + pub fn write(&mut self, path: &Path, content: Box<[u8]>) -> Result<(), LsError> { + let canonical = self.pwd().join(path); + let mut tree = &mut self.data; + + if let Some(parents) = canonical.parent() { + for c in parents { + match tree { + Tree::Directory(d) => { + tree = d + .get_mut(c.to_str().unwrap()) + .ok_or(LsError::NoSuchFileOrDirectory)?; + } + Tree::File(_) => { + return Err(LsError::NotDirectory); + } + } + } + } + + match tree { + Tree::Directory(v) => { + match v.entry( + canonical + .components() + .next_back() + .unwrap() + .as_os_str() + .to_str() + .unwrap() + .to_string(), + ) { + Entry::Vacant(v) => { + v.insert(Box::new(Tree::File(content))); + Ok(()) + } + Entry::Occupied(mut o) if matches!(o.get().as_ref(), Tree::File(_)) => { + o.insert(Box::new(Tree::File(content))); + Ok(()) + } + Entry::Occupied(_) => Err(LsError::IsADirectory), + } + } + Tree::File(_) => Err(LsError::NotDirectory), + } + } + #[allow(clippy::unused_self)] - pub fn ls(&self, _dir: Option<&str>) -> &[&str] { - &[] + pub fn ls<'a>(&'a self, dir: Option<&'a Path>) -> Result, LsError> { + let canonical = if let Some(dir) = dir { + Cow::Owned(self.pwd().join(dir)) + } else { + Cow::Borrowed(self.pwd()) + }; + + let mut tree = &self.data; + + for c in canonical.as_ref() { + match tree { + Tree::Directory(d) => { + tree = d + .get(c.to_str().unwrap()) + .ok_or(LsError::NoSuchFileOrDirectory)?; + } + Tree::File(_) => { + return Err(LsError::NotDirectory); + } + } + } + + match tree { + Tree::Directory(v) => Ok(v.keys().map(String::as_str).collect()), + Tree::File(_) => Ok(vec![dir.unwrap_or(self.pwd()).to_str().unwrap()]), + } + } +} + +#[derive(Debug)] +pub enum LsError { + NotDirectory, + NoSuchFileOrDirectory, + IsADirectory, + FileExists, +} + +impl Display for LsError { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + LsError::NoSuchFileOrDirectory => "No such file or directory", + LsError::NotDirectory => "Not a directory", + LsError::IsADirectory => "Is a directory", + LsError::FileExists => "File exists", + }) } } -- libgit2 1.7.2