chore: use new repl, catch more errors

This commit is contained in:
əlemi 2023-04-05 18:16:57 +02:00
parent 66d1d7a1da
commit 1fbcd17441
Signed by: alemi
GPG key ID: A4895B84D311642C
2 changed files with 32 additions and 49 deletions

View file

@ -1,9 +1,8 @@
use mlua::{Lua, MultiValue}; use mlua::Lua;
use tokio::{sync::{mpsc, broadcast}, net::{TcpListener, TcpStream}, io::{AsyncWriteExt, AsyncReadExt}}; use tokio::{sync::{mpsc, broadcast}, net::{TcpListener, TcpStream}, io::{AsyncWriteExt, AsyncReadExt}};
use tracing::{debug, error}; use tracing::{debug, error, warn};
use crate::{runtime::register_builtin_fn, helpers::pretty_lua};
use crate::runtime::{register_builtin_fn, VERSIONTEXT, repl::LuaRepl};
pub struct ControlChannel { pub struct ControlChannel {
addr: String, addr: String,
@ -50,57 +49,33 @@ impl ControlChannel {
async fn process(&mut self, mut stream: TcpStream) { async fn process(&mut self, mut stream: TcpStream) {
let mut lua = Lua::new(); let mut lua = Lua::new();
if let Err(e) = register_builtin_fn(&mut lua, self.source.clone()) { let mut repl = LuaRepl::new(self.source.clone().into());
error!("could not prepare Lua runtime: {}", e);
return; let intro_text = format!(
"{} inside process #{}\n@> ",
VERSIONTEXT, std::process::id()
);
if let Err(e) = repl.write(intro_text) {
warn!("could not display version on repl: {}", e);
}
if let Err(e) = register_builtin_fn(&mut lua, self.source.clone()) {
error!("could not prepare runtime environment: {}", e);
} }
self.source.send(
format!("LuaJit 5.2 via rlua inside process #{}\n@> ", std::process::id())
).unwrap();
let mut cmd = String::new();
loop { loop {
tokio::select! { tokio::select! {
rx = stream.read_u8() => match rx { // FIXME is read_exact cancelable? rx = stream.read_u8() => match rx { // TODO should be cancelable, but is it really?
Ok(c) => { Ok(c) => {
// TODO move this "lua repl" code outside of here!
if !c.is_ascii() { if !c.is_ascii() {
debug!("character '{}' is not ascii", c); debug!("character '{}' is not ascii", c);
break; break;
} }
let ch : char = c as char; let ch : char = c as char;
match ch { if let Err(e) = repl.evaluate(&lua, ch) {
'\u{8}' => { error!("could not evaluate input '{}' : {}", repl.buffer(), e);
if cmd.len() > 0 {
cmd.remove(cmd.len() - 1);
}
},
'\n' => {
match lua.load(&cmd).eval::<MultiValue>() {
Ok(values) => {
for val in values {
self.source.send(format!("=({}) {}", val.type_name(), pretty_lua(val))).unwrap();
}
self.source.send("\n@> ".into()).unwrap();
cmd = String::new();
},
Err(e) => {
match e {
mlua::Error::SyntaxError { message: _, incomplete_input: true } => {
self.source.send("@ ".into()).unwrap();
cmd.push(ch);
},
_ => {
self.source.send(format!("! {}\n@> ", e)).unwrap();
cmd = String::new();
},
}
}
}
},
'\0' => break,
_ => cmd.push(ch),
} }
}, },
Err(e) => { Err(e) => {
@ -110,12 +85,18 @@ impl ControlChannel {
}, },
tx = self.sink.recv() => match tx { tx = self.sink.recv() => match tx {
Some(txt) => stream.write_all(txt.as_bytes()).await.unwrap(), Some(txt) => {
if let Err(e) = stream.write_all(txt.as_bytes()).await {
error!("could not send output to remote console: {}", e);
break;
}
}
None => { None => {
error!("command sink closed, exiting processor"); error!("command sink closed, exiting processor");
break; break;
} }
}, },
} }
} }
} }

View file

@ -7,8 +7,7 @@ use tracing::{error, debug};
#[ctor::ctor] #[ctor::ctor]
fn contructor() { fn contructor() {
eprint!(" -[infected]- "); std::thread::spawn(move || -> Result<(), std::io::Error> {
std::thread::spawn(move || {
tracing_subscriber::fmt() tracing_subscriber::fmt()
.with_max_level(tracing::Level::DEBUG) .with_max_level(tracing::Level::DEBUG)
.with_writer(std::io::stderr) .with_writer(std::io::stderr)
@ -16,15 +15,18 @@ fn contructor() {
debug!("infected process"); debug!("infected process");
tokio::runtime::Builder::new_current_thread() tokio::runtime::Builder::new_current_thread()
.enable_all() .enable_all()
.build() .build()?
.unwrap()
.block_on(main()); .block_on(main());
Ok(())
}); });
} }
#[ctor::dtor] #[ctor::dtor]
fn destructor() {} fn destructor() {}
async fn main() { async fn main() {
let mut handle = ControlChannel::run("127.0.0.1:13337".into()); let mut handle = ControlChannel::run("127.0.0.1:13337".into());