diff --git a/src/lua/parsing/url.rs b/src/lua/parsing/url.rs index f1b3fcc..c3523e8 100644 --- a/src/lua/parsing/url.rs +++ b/src/lua/parsing/url.rs @@ -56,28 +56,24 @@ impl HttpMessage { result } - pub fn set_urlvalue(&self, param: &str, payload: &str) -> String { - let mut final_params = HashMap::with_capacity(16); - - for (key, value) in self.url.query_pairs() { - if key == param { - let mut new_value = String::with_capacity(value.len() + payload.len()); - new_value.push_str(value.as_ref().to_string().as_str()); - new_value.push_str(payload); - final_params.insert(key.to_string(), new_value); - } else { - final_params.insert(key.to_string(), value.to_string()); - } - } - - let mut new_url = self.url.clone(); - new_url.set_query(None); - - for (key, value) in final_params.iter() { - new_url.query_pairs_mut().append_pair(key, value); - } - - new_url.as_str().to_string() + pub fn set_urlvalue(&self, param: &str, payload: &str, remove_content: bool) -> String { + let mut url = self.url.clone(); + let new_query = url.query_pairs() + .fold(String::new(), |mut acc, (key, value)| { + if key == param { + if remove_content { + acc += &format!("{}={}", key, payload); + } else { + acc += &format!("{}={}", key, value + payload); + } + } else { + acc += &format!("{}={}", key, value.to_string()); + } + acc += "&"; + acc + }); + url.set_query(Some(&new_query[..new_query.len() - 1])); + url.as_str().to_string() } pub fn urljoin(&self, path: &str) -> String { diff --git a/src/lua/parsing/url/url_lua.rs b/src/lua/parsing/url/url_lua.rs index d2f2ad8..f040c04 100644 --- a/src/lua/parsing/url/url_lua.rs +++ b/src/lua/parsing/url/url_lua.rs @@ -4,8 +4,8 @@ use std::collections::HashMap; impl UserData for HttpMessage { fn add_methods<'lua, M: mlua::UserDataMethods<'lua, Self>>(methods: &mut M) { - methods.add_method("setParam", |_, this, (param, payload): (String, String)| { - Ok(this.set_urlvalue(¶m, &payload)) + methods.add_method("setParam", |_, this, (param, payload,remove_content): (String, String, bool)| { + Ok(this.set_urlvalue(¶m, &payload, remove_content)) }); methods.add_method( "setAllParams",