536 lines
19 KiB
Rust
536 lines
19 KiB
Rust
use std::{
|
|
collections::{BTreeSet, VecDeque},
|
|
str::FromStr,
|
|
};
|
|
|
|
use anyhow::bail;
|
|
use gc_arena::{Collect, GcRefLock, Rootable};
|
|
use itertools::Itertools;
|
|
use piccolo::{Callback, Context, IntoValue, Table, UserData, Value};
|
|
use regex::Regex;
|
|
|
|
use crate::{
|
|
lua_engine::frames::try_unwrap_frame,
|
|
parsing::{parse_commands, quote_string, ArgumentGuard, ParsedArgument, ParsedCommand},
|
|
GlobalMemoCell,
|
|
};
|
|
|
|
#[derive(Default, Debug, Collect)]
|
|
#[collect(require_static)]
|
|
pub struct MatchSubTable {
|
|
contents: Vec<MatchRecord>,
|
|
}
|
|
|
|
impl MatchSubTable {
|
|
pub fn to_value<'gc>(&self, ctx: Context<'gc>) -> anyhow::Result<Value<'gc>> {
|
|
let table = Table::new(&ctx);
|
|
for record in self.contents.iter() {
|
|
table.set(
|
|
ctx,
|
|
ctx.intern(record.match_text.as_bytes()),
|
|
ctx.intern(record.sub_text.as_bytes()),
|
|
)?;
|
|
}
|
|
Ok(table.into_value(ctx))
|
|
}
|
|
|
|
pub fn try_sub(&self, input: &str) -> Option<Vec<ParsedCommand>> {
|
|
let cleaned_input = strip_ansi_escapes::strip_str(input);
|
|
for record in self.contents.iter().rev() {
|
|
if let Some(matched) = record.match_regex.captures(&cleaned_input) {
|
|
let vec = Some(
|
|
record
|
|
.sub_commands
|
|
.iter()
|
|
.map(|subcmd| ParsedCommand {
|
|
arguments: subcmd
|
|
.arguments
|
|
.iter()
|
|
.map(|subarg| {
|
|
let unquoted_text = subarg
|
|
.text_parts
|
|
.iter()
|
|
.map(|tp| match tp {
|
|
SubTextPart::Literal(t) => t.as_str(),
|
|
SubTextPart::Variable(v) => {
|
|
if let Ok(v) = <usize as FromStr>::from_str(v) {
|
|
matched.get(v).map_or("", |v| v.as_str())
|
|
} else {
|
|
matched.name(v).map_or("", |v| v.as_str())
|
|
}
|
|
}
|
|
})
|
|
.join("");
|
|
ParsedArgument {
|
|
guard: if subarg.guard.is_none()
|
|
&& unquoted_text.contains(';')
|
|
{
|
|
Some(ArgumentGuard::Paren)
|
|
} else {
|
|
subarg.guard.clone()
|
|
},
|
|
text: unquoted_text.clone(),
|
|
quoted_text: quote_string(&unquoted_text),
|
|
}
|
|
})
|
|
.collect(),
|
|
})
|
|
.collect(),
|
|
);
|
|
return vec;
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
pub fn add_record(&mut self, match_text: &str, sub_text: &str) -> anyhow::Result<()> {
|
|
self.remove_record(match_text).unwrap_or(());
|
|
let rex = Regex::new(match_text)?;
|
|
|
|
let parse_result = parse_commands(sub_text);
|
|
let sub_commands: Vec<SubCommand> = parse_result
|
|
.commands
|
|
.into_iter()
|
|
.map(|cmd| {
|
|
Ok(SubCommand {
|
|
arguments: cmd
|
|
.arguments
|
|
.into_iter()
|
|
.map(parsedarg_to_subarg)
|
|
.collect::<anyhow::Result<VecDeque<SubArgument>>>()?,
|
|
})
|
|
})
|
|
.collect::<anyhow::Result<Vec<SubCommand>>>()?;
|
|
|
|
let vars: BTreeSet<String> = sub_commands
|
|
.iter()
|
|
.flat_map(|sc| {
|
|
sc.arguments.iter().flat_map(|arg| {
|
|
arg.text_parts.iter().filter_map(|tp| match tp {
|
|
SubTextPart::Variable(v) => Some(v.clone()),
|
|
_ => None,
|
|
})
|
|
})
|
|
})
|
|
.collect();
|
|
let max_captures = rex.captures_len();
|
|
let valid_vars: BTreeSet<String> = (0..max_captures)
|
|
.map(|n| n.to_string())
|
|
.chain(rex.capture_names().filter_map(|o| o.map(|n| n.to_string())))
|
|
.collect();
|
|
let invalid_vars: Vec<String> = vars.difference(&valid_vars).cloned().collect();
|
|
if !invalid_vars.is_empty() {
|
|
bail!("Invalid variables in substitution: {:?}", invalid_vars);
|
|
}
|
|
|
|
self.contents.push(MatchRecord {
|
|
match_text: match_text.to_owned(),
|
|
match_regex: rex,
|
|
sub_text: sub_text.to_owned(),
|
|
sub_commands,
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
pub fn remove_record(&mut self, match_text: &str) -> anyhow::Result<()> {
|
|
match self
|
|
.contents
|
|
.iter()
|
|
.enumerate()
|
|
.find(|(_idx, rec)| rec.match_text == match_text)
|
|
{
|
|
None => bail!("No matching record found."),
|
|
Some((idx, _)) => self.contents.remove(idx),
|
|
};
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn parsedarg_to_subarg(parsedarg: ParsedArgument) -> anyhow::Result<SubArgument> {
|
|
let mut text_parts: Vec<SubTextPart> = vec![];
|
|
let mut iter = parsedarg.text.chars().peekable();
|
|
let mut buf = String::new();
|
|
'outer: loop {
|
|
match iter.next() {
|
|
None => break 'outer,
|
|
Some('$') => match iter.peek() {
|
|
None => {
|
|
bail!("substitution ends in $ which is invalid.")
|
|
}
|
|
Some('$') => {
|
|
iter.next();
|
|
buf.push('$')
|
|
}
|
|
Some('{') => {
|
|
iter.next();
|
|
if !buf.is_empty() {
|
|
text_parts.push(SubTextPart::Literal(buf));
|
|
}
|
|
buf = String::new();
|
|
'inner: loop {
|
|
match iter.next() {
|
|
None => {
|
|
bail!("substitution opened with {{ is never closed.")
|
|
}
|
|
Some('}') => {
|
|
if buf.is_empty() {
|
|
bail!("substitution of empty variable name.");
|
|
}
|
|
text_parts.push(SubTextPart::Variable(buf));
|
|
buf = String::new();
|
|
break 'inner;
|
|
}
|
|
Some(c) => buf.push(c),
|
|
}
|
|
}
|
|
}
|
|
Some(_) => {
|
|
if !buf.is_empty() {
|
|
text_parts.push(SubTextPart::Literal(buf));
|
|
}
|
|
buf = String::new();
|
|
'inner: loop {
|
|
match iter.peek() {
|
|
None => {
|
|
text_parts.push(SubTextPart::Variable(buf));
|
|
buf = String::new();
|
|
break 'outer;
|
|
}
|
|
Some(c) if *c == '_' || c.is_ascii_alphanumeric() => {
|
|
buf.push(*c);
|
|
iter.next();
|
|
}
|
|
Some(_) => {
|
|
if buf.is_empty() {
|
|
bail!("substitution of empty variable name.");
|
|
}
|
|
text_parts.push(SubTextPart::Variable(buf));
|
|
buf = String::new();
|
|
break 'inner;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
},
|
|
Some(c) => buf.push(c),
|
|
}
|
|
}
|
|
if !buf.is_empty() {
|
|
text_parts.push(SubTextPart::Literal(buf))
|
|
}
|
|
Ok(SubArgument {
|
|
guard: parsedarg.guard,
|
|
text_parts,
|
|
})
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct MatchRecord {
|
|
pub match_text: String,
|
|
pub match_regex: Regex,
|
|
pub sub_text: String,
|
|
// We parse into into commands and arguments upfront, before substitution.
|
|
// This reduces the risk of security problems.
|
|
pub sub_commands: Vec<SubCommand>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct SubCommand {
|
|
pub arguments: VecDeque<SubArgument>,
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
pub struct SubArgument {
|
|
pub guard: Option<ArgumentGuard>,
|
|
pub text_parts: Vec<SubTextPart>,
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
pub enum SubTextPart {
|
|
Literal(String),
|
|
Variable(String),
|
|
}
|
|
|
|
pub fn create_match_table<'gc, 'a>(ctx: Context<'gc>) -> Callback<'gc> {
|
|
Callback::from_fn(&ctx, move |ctx, _ex, mut stack| {
|
|
let _: () = stack.consume(ctx)?;
|
|
let user_data = UserData::<'gc>::new::<Rootable!['gcb => GcRefLock<'gcb, MatchSubTable>]>(
|
|
&ctx,
|
|
GcRefLock::new(&ctx, <MatchSubTable as Default>::default().into()),
|
|
);
|
|
let match_table_class: Table = ctx
|
|
.get_global::<Table>("classes")?
|
|
.get(ctx, "match_table")?;
|
|
user_data.set_metatable(&ctx, Some(match_table_class));
|
|
|
|
stack.push_back(user_data.into_value(ctx));
|
|
Ok(piccolo::CallbackReturn::Return)
|
|
})
|
|
}
|
|
|
|
pub fn match_table_add<'gc, 'a>(ctx: Context<'gc>) -> Callback<'gc> {
|
|
Callback::from_fn(&ctx, move |ctx, _ex, mut stack| {
|
|
let (match_table, match_text, sub_text): (UserData, piccolo::String, piccolo::String) =
|
|
stack.consume(ctx)?;
|
|
match_table
|
|
.downcast::<Rootable!['gcb => GcRefLock<'gcb, MatchSubTable>]>()?
|
|
.borrow_mut(&ctx)
|
|
.add_record(match_text.to_str()?, sub_text.to_str()?)?;
|
|
Ok(piccolo::CallbackReturn::Return)
|
|
})
|
|
}
|
|
|
|
pub fn match_table_remove<'gc, 'a>(ctx: Context<'gc>) -> Callback<'gc> {
|
|
Callback::from_fn(&ctx, move |ctx, _ex, mut stack| {
|
|
let (match_table, match_text): (UserData, piccolo::String) = stack.consume(ctx)?;
|
|
match_table
|
|
.downcast::<Rootable!['gcb => GcRefLock<'gcb, MatchSubTable>]>()?
|
|
.borrow_mut(&ctx)
|
|
.remove_record(match_text.to_str()?)?;
|
|
Ok(piccolo::CallbackReturn::Return)
|
|
})
|
|
}
|
|
|
|
pub fn match_table_lua_table<'gc, 'a>(ctx: Context<'gc>) -> Callback<'gc> {
|
|
Callback::from_fn(&ctx, move |ctx, _ex, mut stack| {
|
|
let match_table: UserData = stack.consume(ctx)?;
|
|
stack.push_back(
|
|
match_table
|
|
.downcast::<Rootable!['gcb => GcRefLock<'gcb, MatchSubTable>]>()?
|
|
.borrow_mut(&ctx)
|
|
.to_value(ctx)?,
|
|
);
|
|
Ok(piccolo::CallbackReturn::Return)
|
|
})
|
|
}
|
|
|
|
pub fn match_table_try_run_sub<'gc, 'a>(
|
|
ctx: Context<'gc>,
|
|
global_memo: &'a GlobalMemoCell,
|
|
) -> Callback<'gc> {
|
|
let global_memo = global_memo.clone();
|
|
Callback::from_fn(&ctx, move |ctx, _ex, mut stack| {
|
|
let (match_table, sub, frame): (UserData, piccolo::String, Value) = stack.consume(ctx)?;
|
|
let frame = try_unwrap_frame(ctx, &frame)?;
|
|
|
|
let cmds = match_table
|
|
.downcast::<Rootable!['gcb => GcRefLock<'gcb, MatchSubTable>]>()?
|
|
.borrow()
|
|
.try_sub(sub.to_str()?);
|
|
|
|
match cmds {
|
|
None => stack.push_back(false.into_value(ctx)),
|
|
Some(cmds) => {
|
|
let mut cq = global_memo.command_queue.borrow_mut();
|
|
for cmd in cmds.into_iter().rev() {
|
|
cq.push_front((frame.clone(), cmd));
|
|
}
|
|
stack.push_back(Value::Boolean(true))
|
|
}
|
|
}
|
|
Ok(piccolo::CallbackReturn::Return)
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn parsedarg_to_subarg_works() {
|
|
assert_eq!(
|
|
parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "hello world!".to_owned(),
|
|
quoted_text: "hello world!".to_owned()
|
|
})
|
|
.unwrap(),
|
|
SubArgument {
|
|
guard: None,
|
|
text_parts: vec![SubTextPart::Literal("hello world!".to_owned())]
|
|
}
|
|
);
|
|
assert_eq!(
|
|
parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "hello $adjective ${my world}".to_owned(),
|
|
quoted_text: "hello $adjective ${my world}".to_owned()
|
|
})
|
|
.unwrap(),
|
|
SubArgument {
|
|
guard: None,
|
|
text_parts: vec![
|
|
SubTextPart::Literal("hello ".to_owned()),
|
|
SubTextPart::Variable("adjective".to_owned()),
|
|
SubTextPart::Literal(" ".to_owned()),
|
|
SubTextPart::Variable("my world".to_owned()),
|
|
]
|
|
}
|
|
);
|
|
assert_eq!(
|
|
parsedarg_to_subarg(ParsedArgument {
|
|
guard: Some(ArgumentGuard::DoubleQuote),
|
|
text: "hello $adjective${my world}${your world} end".to_owned(),
|
|
quoted_text: "hello $adjective$\\{my world\\}$\\{your world\\} end".to_owned()
|
|
})
|
|
.unwrap(),
|
|
SubArgument {
|
|
guard: Some(ArgumentGuard::DoubleQuote),
|
|
text_parts: vec![
|
|
SubTextPart::Literal("hello ".to_owned()),
|
|
SubTextPart::Variable("adjective".to_owned()),
|
|
SubTextPart::Variable("my world".to_owned()),
|
|
SubTextPart::Variable("your world".to_owned()),
|
|
SubTextPart::Literal(" end".to_owned()),
|
|
]
|
|
}
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn parsedarg_rejects_invalid() {
|
|
assert!(parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "${untermin".to_owned(),
|
|
quoted_text: "${untermin".to_owned()
|
|
})
|
|
.is_err());
|
|
assert!(parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "$foo$".to_owned(),
|
|
quoted_text: "$foo$".to_owned()
|
|
})
|
|
.is_err());
|
|
assert!(parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "$ hello".to_owned(),
|
|
quoted_text: "$ hello".to_owned()
|
|
})
|
|
.is_err());
|
|
assert!(parsedarg_to_subarg(ParsedArgument {
|
|
guard: None,
|
|
text: "My name is ${}".to_owned(),
|
|
quoted_text: "My name is ${}".to_owned()
|
|
})
|
|
.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn matchsubtable_works() {
|
|
let mut table: MatchSubTable = Default::default();
|
|
table
|
|
.add_record(
|
|
"^foo (?<bar>[a-z]+) baz",
|
|
"\\\"Someone is talking $bar about foo baz?;:flexes his ${bar}",
|
|
)
|
|
.expect("adding record failed");
|
|
assert_eq!(table.try_sub("unrelated babble"), None);
|
|
assert_eq!(
|
|
table.try_sub("foo woots baz\r\n"),
|
|
Some(vec![
|
|
ParsedCommand {
|
|
arguments: [
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "\"Someone".to_owned(),
|
|
quoted_text: "\\\"Someone".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "is".to_owned(),
|
|
quoted_text: "is".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "talking".to_owned(),
|
|
quoted_text: "talking".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "woots".to_owned(),
|
|
quoted_text: "woots".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "about".to_owned(),
|
|
quoted_text: "about".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "foo".to_owned(),
|
|
quoted_text: "foo".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "baz?".to_owned(),
|
|
quoted_text: "baz?".to_owned()
|
|
}
|
|
]
|
|
.into()
|
|
},
|
|
ParsedCommand {
|
|
arguments: [
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: ":flexes".to_owned(),
|
|
quoted_text: ":flexes".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "his".to_owned(),
|
|
quoted_text: "his".to_owned()
|
|
},
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "woots".to_owned(),
|
|
quoted_text: "woots".to_owned()
|
|
}
|
|
]
|
|
.into()
|
|
}
|
|
])
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn matchsubtable_resists_command_injection() {
|
|
let mut table: MatchSubTable = Default::default();
|
|
table
|
|
.add_record("^foo (.*)", "safe_command $1")
|
|
.expect("adding record failed");
|
|
let result = table
|
|
.try_sub("foo pwned!};dangerous_command {")
|
|
.expect("didn't match");
|
|
let expected = ParsedCommand {
|
|
arguments: [
|
|
ParsedArgument {
|
|
guard: None,
|
|
text: "safe_command".to_owned(),
|
|
quoted_text: "safe_command".to_owned(),
|
|
},
|
|
ParsedArgument {
|
|
guard: Some(ArgumentGuard::Paren),
|
|
text: "pwned!};dangerous_command {".to_owned(),
|
|
quoted_text: "pwned!\\};dangerous_command \\{".to_owned(),
|
|
},
|
|
]
|
|
.into(),
|
|
};
|
|
assert_eq!(result, vec![expected.clone()]);
|
|
let ser_result = result[0].to_string();
|
|
assert_eq!(
|
|
ser_result,
|
|
"safe_command {pwned!\\};dangerous_command \\{}".to_owned()
|
|
);
|
|
assert_eq!(parse_commands(&ser_result).commands, vec![expected]);
|
|
}
|
|
|
|
#[test]
|
|
fn matchsubtable_rejects_invalid() {
|
|
let mut table: MatchSubTable = Default::default();
|
|
assert!(table
|
|
.add_record("^foo (?<bar>[a-z]+) baz", "$wrong")
|
|
.is_err())
|
|
}
|
|
}
|