diff --git a/lib/src/fnc/string.rs b/lib/src/fnc/string.rs index 4e9a996a..f14ab722 100644 --- a/lib/src/fnc/string.rs +++ b/lib/src/fnc/string.rs @@ -2,8 +2,23 @@ use crate::err::Error; use crate::fnc::util::string; use crate::sql::value::Value; +/// Returns `true` if a string of this length is too much to allocate. +fn limit(name: &str, n: usize) -> Result<(), Error> { + const LIMIT: usize = 2usize.pow(20); + if n > LIMIT { + Err(Error::InvalidArguments { + name: name.to_owned(), + message: format!("Output must not exceed {LIMIT} bytes."), + }) + } else { + Ok(()) + } +} + pub fn concat(args: Vec) -> Result { - Ok(args.into_iter().map(|x| x.as_string()).collect::>().concat().into()) + let strings = args.into_iter().map(Value::as_string).collect::>(); + limit("string::concat", strings.iter().map(String::len).sum::())?; + Ok(strings.concat().into()) } pub fn contains((val, check): (String, String)) -> Result { @@ -20,10 +35,19 @@ pub fn join(args: Vec) -> Result { name: String::from("string::join"), message: String::from("Expected at least one argument"), })?; + + let strings = args.collect::>(); + limit( + "string::join", + strings + .len() + .saturating_mul(chr.len()) + .saturating_add(strings.iter().map(String::len).sum::()), + )?; + // FIXME: Use intersperse to avoid intermediate allocation once stable // https://github.com/rust-lang/rust/issues/79524 - let val = args.collect::>().join(&chr); - Ok(val.into()) + Ok(strings.join(&chr).into()) } pub fn len((string,): (String,)) -> Result { @@ -36,18 +60,18 @@ pub fn lowercase((string,): (String,)) -> Result { } pub fn repeat((val, num): (String, usize)) -> Result { - const LIMIT: usize = 2usize.pow(20); - if val.len().saturating_mul(num) > LIMIT { - Err(Error::InvalidArguments { - name: String::from("string::repeat"), - message: format!("Output must not exceed {LIMIT} bytes."), - }) - } else { - Ok(val.repeat(num).into()) - } + limit("string::repeat", val.len().saturating_mul(num))?; + Ok(val.repeat(num).into()) } pub fn replace((val, old, new): (String, String, String)) -> Result { + if new.len() > old.len() { + let increase = new.len() - old.len(); + limit( + "string::replace", + val.len().saturating_add(val.matches(&old).count().saturating_mul(increase)), + )?; + } Ok(val.replace(&old, &new).into()) }