diff --git a/rust/src/dns/log.rs b/rust/src/dns/log.rs index 3de67c9bc75b..01dac80d0b83 100644 --- a/rust/src/dns/log.rs +++ b/rust/src/dns/log.rs @@ -399,7 +399,7 @@ pub fn dns_print_addr(addr: &Vec) -> std::string::String { /// Log SOA section fields. fn dns_log_soa(soa: &DNSRDataSOA) -> Result { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object()?; js.set_string_from_bytes("mname", &soa.mname)?; js.set_string_from_bytes("rname", &soa.rname)?; @@ -415,7 +415,7 @@ fn dns_log_soa(soa: &DNSRDataSOA) -> Result { /// Log SSHFP section fields. fn dns_log_sshfp(sshfp: &DNSRDataSSHFP) -> Result { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object()?; let mut hex = Vec::new(); for byte in &sshfp.fingerprint { @@ -432,7 +432,7 @@ fn dns_log_sshfp(sshfp: &DNSRDataSSHFP) -> Result { /// Log SRV section fields. fn dns_log_srv(srv: &DNSRDataSRV) -> Result { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object()?; js.set_uint("priority", srv.priority as u64)?; js.set_uint("weight", srv.weight as u64)?; @@ -444,7 +444,7 @@ fn dns_log_srv(srv: &DNSRDataSRV) -> Result { } fn dns_log_json_answer_detail(answer: &DNSAnswerEntry) -> Result { - let mut jsa = JsonBuilder::new_object(); + let mut jsa = JsonBuilder::try_new_object()?; jsa.set_string_from_bytes("rrname", &answer.name)?; jsa.set_string("rrtype", &dns_rrtype_string(answer.rrtype))?; @@ -516,7 +516,7 @@ fn dns_log_json_answer( js.set_string("rcode", &dns_rcode_string(header.flags))?; if !response.answers.is_empty() { - let mut js_answers = JsonBuilder::new_array(); + let mut js_answers = JsonBuilder::try_new_array()?; // For grouped answers we use a HashMap keyed by the rrtype. let mut answer_types = HashMap::new(); @@ -527,7 +527,7 @@ fn dns_log_json_answer( match &answer.data { DNSRData::A(addr) | DNSRData::AAAA(addr) => { if !answer_types.contains_key(&type_string) { - answer_types.insert(type_string.to_string(), JsonBuilder::new_array()); + answer_types.insert(type_string.to_string(), JsonBuilder::try_new_array()?); } if let Some(a) = answer_types.get_mut(&type_string) { a.append_string(&dns_print_addr(addr))?; @@ -540,7 +540,7 @@ fn dns_log_json_answer( | DNSRData::NULL(bytes) | DNSRData::PTR(bytes) => { if !answer_types.contains_key(&type_string) { - answer_types.insert(type_string.to_string(), JsonBuilder::new_array()); + answer_types.insert(type_string.to_string(), JsonBuilder::try_new_array()?); } if let Some(a) = answer_types.get_mut(&type_string) { a.append_string_from_bytes(bytes)?; @@ -548,7 +548,7 @@ fn dns_log_json_answer( } DNSRData::SOA(soa) => { if !answer_types.contains_key(&type_string) { - answer_types.insert(type_string.to_string(), JsonBuilder::new_array()); + answer_types.insert(type_string.to_string(), JsonBuilder::try_new_array()?); } if let Some(a) = answer_types.get_mut(&type_string) { a.append_object(&dns_log_soa(soa)?)?; @@ -556,7 +556,7 @@ fn dns_log_json_answer( } DNSRData::SSHFP(sshfp) => { if !answer_types.contains_key(&type_string) { - answer_types.insert(type_string.to_string(), JsonBuilder::new_array()); + answer_types.insert(type_string.to_string(), JsonBuilder::try_new_array()?); } if let Some(a) = answer_types.get_mut(&type_string) { a.append_object(&dns_log_sshfp(sshfp)?)?; @@ -564,7 +564,7 @@ fn dns_log_json_answer( } DNSRData::SRV(srv) => { if !answer_types.contains_key(&type_string) { - answer_types.insert(type_string.to_string(), JsonBuilder::new_array()); + answer_types.insert(type_string.to_string(), JsonBuilder::try_new_array()?); } if let Some(a) = answer_types.get_mut(&type_string) { a.append_object(&dns_log_srv(srv)?)?; diff --git a/rust/src/jsonbuilder.rs b/rust/src/jsonbuilder.rs index abf579aae35b..c1f466563be7 100644 --- a/rust/src/jsonbuilder.rs +++ b/rust/src/jsonbuilder.rs @@ -17,6 +17,8 @@ #![allow(clippy::missing_safety_doc)] +use std::cmp::max; +use std::collections::TryReserveError; use std::ffi::CStr; use std::os::raw::c_char; use std::str::Utf8Error; @@ -27,6 +29,7 @@ const INIT_SIZE: usize = 4096; pub enum JsonError { InvalidState, Utf8Error(Utf8Error), + Memory, } impl std::error::Error for JsonError {} @@ -36,10 +39,17 @@ impl std::fmt::Display for JsonError { match self { JsonError::InvalidState => write!(f, "invalid state"), JsonError::Utf8Error(ref e) => e.fmt(f), + JsonError::Memory => write!(f, "memory error"), } } } +impl From for JsonError { + fn from(_: TryReserveError) -> Self { + JsonError::Memory + } +} + impl From for JsonError { fn from(e: Utf8Error) -> Self { JsonError::Utf8Error(e) @@ -98,47 +108,78 @@ pub struct JsonBuilder { impl JsonBuilder { /// Returns a new JsonBuilder in object state. - pub fn new_object() -> Self { - Self::new_object_with_capacity(INIT_SIZE) + pub fn try_new_object() -> Result { + Self::try_new_object_with_capacity(INIT_SIZE) } - pub fn new_object_with_capacity(capacity: usize) -> Self { - let mut buf = String::with_capacity(capacity); + pub fn try_new_object_with_capacity(capacity: usize) -> Result { + let mut buf = String::new(); + buf.try_reserve(capacity)?; buf.push('{'); - Self { + let mut state = Vec::new(); + state.try_reserve(32)?; + state.extend_from_slice(&[State::None, State::ObjectFirst]); + Ok(Self { buf, - state: vec![State::None, State::ObjectFirst], + state, init_type: Type::Object, - } + }) } /// Returns a new JsonBuilder in array state. - pub fn new_array() -> Self { - Self::new_array_with_capacity(INIT_SIZE) + pub fn try_new_array() -> Result { + Self::try_new_array_with_capacity(INIT_SIZE) } - pub fn new_array_with_capacity(capacity: usize) -> Self { - let mut buf = String::with_capacity(capacity); + pub fn try_new_array_with_capacity(capacity: usize) -> Result { + let mut buf = String::new(); + buf.try_reserve(capacity)?; buf.push('['); - Self { + let mut state = Vec::new(); + state.try_reserve(32)?; + state.extend_from_slice(&[State::None, State::ArrayFirst]); + Ok(Self { buf, - state: vec![State::None, State::ArrayFirst], + state, init_type: Type::Array, + }) + } + + /// A wrapper around String::push that pre-allocates data return + /// an error if unable to. + pub fn push(&mut self, ch: char) -> Result<&mut Self, JsonError> { + if self.buf.capacity() == self.buf.len() { + self.buf.try_reserve(INIT_SIZE)?; } + self.buf.push(ch); + Ok(self) + } + + /// A wrapper around String::push_str that pre-allocates data + /// return an error if unable to. + pub fn push_str(&mut self, s: &str) -> Result<&mut Self, JsonError> { + if self.buf.capacity() < self.buf.len() + s.len() { + self.buf.try_reserve(max(INIT_SIZE, s.len()))?; + } + self.buf.push_str(s); + Ok(self) } // Reset the builder to its initial state, without losing // the current capacity. pub fn reset(&mut self) { self.buf.truncate(0); + self.state.clear(); match self.init_type { Type::Array => { self.buf.push('['); - self.state = vec![State::None, State::ArrayFirst]; + self.state + .extend_from_slice(&[State::None, State::ArrayFirst]); } Type::Object => { self.buf.push('{'); - self.state = vec![State::None, State::ObjectFirst]; + self.state + .extend_from_slice(&[State::None, State::ObjectFirst]); } } } @@ -147,12 +188,12 @@ impl JsonBuilder { pub fn close(&mut self) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectFirst | State::ObjectNth => { - self.buf.push('}'); + self.push('}')?; self.pop_state(); Ok(self) } State::ArrayFirst | State::ArrayNth => { - self.buf.push(']'); + self.push(']')?; self.pop_state(); Ok(self) } @@ -173,8 +214,12 @@ impl JsonBuilder { } /// Move to a new state. - fn push_state(&mut self, state: State) { + fn push_state(&mut self, state: State) -> Result<(), JsonError> { + if self.state.len() == self.state.capacity() { + self.state.try_reserve(32)?; + } self.state.push(state); + Ok(()) } /// Go back to the previous state. @@ -214,20 +259,20 @@ impl JsonBuilder { pub fn open_object(&mut self, key: &str) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectFirst => { - self.buf.push('"'); + self.push('"')?; self.set_state(State::ObjectNth); } State::ObjectNth => { - self.buf.push_str(",\""); + self.push_str(",\"")?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push_str(key); - self.buf.push_str("\":{"); - self.push_state(State::ObjectFirst); + self.push_str(key)?; + self.push_str("\":{")?; + self.push_state(State::ObjectFirst)?; Ok(self) } @@ -240,16 +285,16 @@ impl JsonBuilder { match self.current_state() { State::ArrayFirst => {} State::ArrayNth => { - self.buf.push(','); + self.push(',')?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push('{'); + self.push('{')?; self.set_state(State::ArrayNth); - self.push_state(State::ObjectFirst); + self.push_state(State::ObjectFirst)?; Ok(self) } @@ -262,18 +307,18 @@ impl JsonBuilder { match self.current_state() { State::ObjectFirst => {} State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":["); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":[")?; self.set_state(State::ObjectNth); - self.push_state(State::ArrayFirst); + self.push_state(State::ArrayFirst)?; Ok(self) } @@ -286,7 +331,7 @@ impl JsonBuilder { Ok(self) } State::ArrayNth => { - self.buf.push(','); + self.push(',')?; self.encode_string(val)?; Ok(self) } @@ -300,7 +345,7 @@ impl JsonBuilder { pub fn append_string_from_bytes(&mut self, val: &[u8]) -> Result<&mut Self, JsonError> { match std::str::from_utf8(val) { Ok(s) => self.append_string(s), - Err(_) => self.append_string(&string_from_bytes(val)), + Err(_) => self.append_string(&try_string_from_bytes(val)?), } } @@ -308,17 +353,17 @@ impl JsonBuilder { pub fn append_base64(&mut self, val: &[u8]) -> Result<&mut Self, JsonError> { match self.current_state() { State::ArrayFirst => { - self.buf.push('"'); - base64::encode_config_buf(val, base64::STANDARD, &mut self.buf); - self.buf.push('"'); + self.push('"')?; + self.encode_base64(val)?; + self.push('"')?; self.set_state(State::ArrayNth); Ok(self) } State::ArrayNth => { - self.buf.push(','); - self.buf.push('"'); - base64::encode_config_buf(val, base64::STANDARD, &mut self.buf); - self.buf.push('"'); + self.push(',')?; + self.push('"')?; + self.encode_base64(val)?; + self.push('"')?; Ok(self) } _ => { @@ -332,23 +377,23 @@ impl JsonBuilder { pub fn append_hex(&mut self, val: &[u8]) -> Result<&mut Self, JsonError> { match self.current_state() { State::ArrayFirst => { - self.buf.push('"'); + self.push('"')?; for i in 0..val.len() { - self.buf.push(HEX[(val[i] >> 4) as usize] as char); - self.buf.push(HEX[(val[i] & 0xf) as usize] as char); + self.push(HEX[(val[i] >> 4) as usize] as char)?; + self.push(HEX[(val[i] & 0xf) as usize] as char)?; } - self.buf.push('"'); + self.push('"')?; self.set_state(State::ArrayNth); Ok(self) } State::ArrayNth => { - self.buf.push(','); - self.buf.push('"'); + self.push(',')?; + self.push('"')?; for i in 0..val.len() { - self.buf.push(HEX[(val[i] >> 4) as usize] as char); - self.buf.push(HEX[(val[i] & 0xf) as usize] as char); + self.push(HEX[(val[i] >> 4) as usize] as char)?; + self.push(HEX[(val[i] & 0xf) as usize] as char)?; } - self.buf.push('"'); + self.push('"')?; Ok(self) } _ => { @@ -365,15 +410,14 @@ impl JsonBuilder { self.set_state(State::ArrayNth); } State::ArrayNth => { - self.buf.push(','); + self.push(',')?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push_str(&val.to_string()); - Ok(self) + self.push_str(&val.to_string()) } pub fn append_float(&mut self, val: f64) -> Result<&mut Self, JsonError> { @@ -382,21 +426,21 @@ impl JsonBuilder { self.set_state(State::ArrayNth); } State::ArrayNth => { - self.buf.push(','); + self.push(',')?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push_str(&val.to_string()); + self.push_str(&val.to_string())?; Ok(self) } pub fn set_object(&mut self, key: &str, js: &JsonBuilder) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -406,10 +450,10 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":"); - self.buf.push_str(&js.buf); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":")?; + self.push_str(&js.buf)?; Ok(self) } @@ -423,14 +467,14 @@ impl JsonBuilder { self.set_state(State::ArrayNth); } State::ArrayNth => { - self.buf.push(','); + self.push(',')?; } _ => { debug_validate_fail!("invalid state"); return Err(JsonError::InvalidState); } } - self.buf.push_str(&js.buf); + self.push_str(&js.buf)?; Ok(self) } @@ -439,7 +483,7 @@ impl JsonBuilder { pub fn set_string(&mut self, key: &str, val: &str) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -449,9 +493,9 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":"); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":")?; self.encode_string(val)?; Ok(self) } @@ -459,7 +503,7 @@ impl JsonBuilder { pub fn set_formatted(&mut self, formatted: &str) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -469,7 +513,7 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push_str(formatted); + self.push_str(formatted)?; Ok(self) } @@ -477,7 +521,7 @@ impl JsonBuilder { pub fn set_string_from_bytes(&mut self, key: &str, val: &[u8]) -> Result<&mut Self, JsonError> { match std::str::from_utf8(val) { Ok(s) => self.set_string(key, s), - Err(_) => self.set_string(key, &string_from_bytes(val)), + Err(_) => self.set_string(key, &try_string_from_bytes(val)?), } } @@ -485,7 +529,7 @@ impl JsonBuilder { pub fn set_base64(&mut self, key: &str, val: &[u8]) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -495,11 +539,11 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":\""); - base64::encode_config_buf(val, base64::STANDARD, &mut self.buf); - self.buf.push('"'); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":\"")?; + self.encode_base64(val)?; + self.push('"')?; Ok(self) } @@ -508,7 +552,7 @@ impl JsonBuilder { pub fn set_hex(&mut self, key: &str, val: &[u8]) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -518,14 +562,14 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":\""); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":\"")?; for i in 0..val.len() { - self.buf.push(HEX[(val[i] >> 4) as usize] as char); - self.buf.push(HEX[(val[i] & 0xf) as usize] as char); + self.push(HEX[(val[i] >> 4) as usize] as char)?; + self.push(HEX[(val[i] & 0xf) as usize] as char)?; } - self.buf.push('"'); + self.push('"')?; Ok(self) } @@ -534,7 +578,7 @@ impl JsonBuilder { pub fn set_uint(&mut self, key: &str, val: u64) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -544,17 +588,17 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":"); - self.buf.push_str(&val.to_string()); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":")?; + self.push_str(&val.to_string())?; Ok(self) } pub fn set_float(&mut self, key: &str, val: f64) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -564,17 +608,17 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); - self.buf.push_str("\":"); - self.buf.push_str(&val.to_string()); + self.push('"')?; + self.push_str(key)?; + self.push_str("\":")?; + self.push_str(&val.to_string())?; Ok(self) } pub fn set_bool(&mut self, key: &str, val: bool) -> Result<&mut Self, JsonError> { match self.current_state() { State::ObjectNth => { - self.buf.push(','); + self.push(',')?; } State::ObjectFirst => { self.set_state(State::ObjectNth); @@ -584,12 +628,12 @@ impl JsonBuilder { return Err(JsonError::InvalidState); } } - self.buf.push('"'); - self.buf.push_str(key); + self.push('"')?; + self.push_str(key)?; if val { - self.buf.push_str("\":true"); + self.push_str("\":true")?; } else { - self.buf.push_str("\":false"); + self.push_str("\":false")?; } Ok(self) } @@ -602,17 +646,27 @@ impl JsonBuilder { /// /// The string is encoded into an intermediate vector as its faster /// than building onto the buffer. + /// + /// TODO: Revisit this, would be nice to build directly onto the + /// existing buffer. #[inline(always)] fn encode_string(&mut self, val: &str) -> Result<(), JsonError> { - let mut buf = vec![0; val.len() * 2 + 2]; + let mut buf = Vec::new(); + + // Start by allocating a reasonable size buffer, it will be + // grown if needed. + buf.try_reserve(val.len() * 2 + 2)?; + buf.resize(val.len() * 2 + 2, 0); + let mut offset = 0; let bytes = val.as_bytes(); buf[offset] = b'"'; offset += 1; for &x in bytes.iter() { if offset + 7 >= buf.capacity() { - let mut extend = vec![0; buf.capacity()]; - buf.append(&mut extend); + // We could be smarter, but just double the buffer size. + buf.try_reserve(buf.capacity())?; + buf.resize(buf.capacity(), 0); } let escape = ESCAPED[x as usize]; if escape == 0 { @@ -642,7 +696,7 @@ impl JsonBuilder { offset += 1; match std::str::from_utf8(&buf[0..offset]) { Ok(s) => { - self.buf.push_str(s); + self.push_str(s)?; } Err(err) => { let error = format!( @@ -651,18 +705,32 @@ impl JsonBuilder { &buf[0..offset], val.as_bytes(), ); - self.buf.push_str(&error); + self.push_str(&error)?; } } Ok(()) } + + fn encode_base64(&mut self, val: &[u8]) -> Result<&mut Self, JsonError> { + let encoded_len = 4 * ((val.len() + 2) / 3); + if self.buf.capacity() < self.buf.len() + encoded_len { + self.buf.try_reserve(encoded_len)?; + } + base64::encode_config_buf(val, base64::STANDARD, &mut self.buf); + Ok(self) + } } /// A Suricata specific function to create a string from bytes when UTF-8 decoding fails. /// /// For bytes over 0x0f, we encode as hex like "\xf2". -fn string_from_bytes(input: &[u8]) -> String { - let mut out = String::with_capacity(input.len()); +fn try_string_from_bytes(input: &[u8]) -> Result { + let mut out = String::new(); + + // Allocate enough data to handle the worst case scenario of every + // byte needing to be presented as a byte. + out.try_reserve(input.len() * 4)?; + for b in input.iter() { if *b < 128 { out.push(*b as char); @@ -670,19 +738,29 @@ fn string_from_bytes(input: &[u8]) -> String { out.push_str(&format!("\\x{:02x}", *b)); } } - return out; + return Ok(out); } #[no_mangle] pub extern "C" fn jb_new_object() -> *mut JsonBuilder { - let boxed = Box::new(JsonBuilder::new_object()); - Box::into_raw(boxed) + match JsonBuilder::try_new_object() { + Ok(js) => { + let boxed = Box::new(js); + Box::into_raw(boxed) + } + Err(_) => std::ptr::null_mut(), + } } #[no_mangle] pub extern "C" fn jb_new_array() -> *mut JsonBuilder { - let boxed = Box::new(JsonBuilder::new_array()); - Box::into_raw(boxed) + match JsonBuilder::try_new_array() { + Ok(js) => { + let boxed = Box::new(js); + Box::into_raw(boxed) + } + Err(_) => std::ptr::null_mut(), + } } #[no_mangle] @@ -908,15 +986,27 @@ pub unsafe extern "C" fn jb_restore_mark(js: &mut JsonBuilder, mark: &mut JsonBu mod test { use super::*; + #[test] + fn test_try_reserve() { + // Just a sanity check that try_reserve works as I expect. + let mut buf = String::new(); + assert_eq!(buf.len(), 0); + assert_eq!(buf.capacity(), 0); + + buf.try_reserve(1).unwrap(); + assert_eq!(buf.len(), 0); + assert!(buf.capacity() >= 1); + } + #[test] fn test_set_bool() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_bool("first", true).unwrap(); assert_eq!(jb.buf, r#"{"first":true"#); jb.set_bool("second", false).unwrap(); assert_eq!(jb.buf, r#"{"first":true,"second":false"#); - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_bool("first", false).unwrap(); assert_eq!(jb.buf, r#"{"first":false"#); jb.set_bool("second", true).unwrap(); @@ -925,7 +1015,7 @@ mod test { #[test] fn test_object_in_object() -> Result<(), JsonError> { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object().unwrap(); js.open_object("object")?; assert_eq!(js.current_state(), State::ObjectFirst); @@ -948,7 +1038,7 @@ mod test { #[test] fn test_empty_array_in_object() -> Result<(), JsonError> { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object().unwrap(); js.open_array("array")?; assert_eq!(js.current_state(), State::ArrayFirst); @@ -966,7 +1056,7 @@ mod test { #[test] #[cfg(not(feature = "debug-validate"))] fn test_array_in_object() -> Result<(), JsonError> { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object().unwrap(); // Attempt to add an item, should fail. assert_eq!( @@ -1002,7 +1092,7 @@ mod test { #[test] fn basic_test() -> Result<(), JsonError> { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object().unwrap(); assert_eq!(js.current_state(), State::ObjectFirst); assert_eq!(js.buf, "{"); @@ -1023,11 +1113,11 @@ mod test { #[test] fn test_combine() -> Result<(), JsonError> { - let mut main = JsonBuilder::new_object(); - let mut obj = JsonBuilder::new_object(); + let mut main = JsonBuilder::try_new_object().unwrap(); + let mut obj = JsonBuilder::try_new_object().unwrap(); obj.close()?; - let mut array = JsonBuilder::new_array(); + let mut array = JsonBuilder::try_new_array().unwrap(); array.append_string("one")?; array.append_uint(2)?; array.close()?; @@ -1042,7 +1132,7 @@ mod test { #[test] fn test_objects_in_array() -> Result<(), JsonError> { - let mut js = JsonBuilder::new_array(); + let mut js = JsonBuilder::try_new_array()?; assert_eq!(js.buf, r#"["#); js.start_object()?; @@ -1071,16 +1161,25 @@ mod test { #[test] fn test_grow() -> Result<(), JsonError> { - let mut jb = JsonBuilder::new_object_with_capacity(1); - assert_eq!(jb.capacity(), 1); - jb.set_string("foo", "bar")?; - assert!(jb.capacity() > 1); + let mut jb = JsonBuilder::try_new_object_with_capacity(1).unwrap(); + + // For efficiency reasons, more capacity may be allocated than + // requested. + assert!(jb.capacity() > 0); + let capacity = jb.capacity(); + + let mut value = String::new(); + for i in 0..capacity { + value.push_str((i % 10).to_string().as_str()); + } + jb.set_string("foo", &value)?; + assert!(jb.capacity() > capacity); Ok(()) } #[test] fn test_reset() -> Result<(), JsonError> { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); assert_eq!(jb.buf, "{"); jb.set_string("foo", "bar")?; let cap = jb.capacity(); @@ -1092,13 +1191,13 @@ mod test { #[test] fn test_append_string_from_bytes() -> Result<(), JsonError> { - let mut jb = JsonBuilder::new_array(); + let mut jb = JsonBuilder::try_new_array().unwrap(); let s = &[0x41, 0x41, 0x41, 0x00]; jb.append_string_from_bytes(s)?; assert_eq!(jb.buf, r#"["AAA\u0000""#); let s = &[0x00, 0x01, 0x02, 0x03]; - let mut jb = JsonBuilder::new_array(); + let mut jb = JsonBuilder::try_new_array().unwrap(); jb.append_string_from_bytes(s)?; assert_eq!(jb.buf, r#"["\u0000\u0001\u0002\u0003""#); @@ -1107,7 +1206,7 @@ mod test { #[test] fn test_set_string_from_bytes() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_string_from_bytes("first", &[]).unwrap(); assert_eq!(jb.buf, r#"{"first":"""#); jb.set_string_from_bytes("second", &[]).unwrap(); @@ -1117,15 +1216,13 @@ mod test { #[test] fn test_append_string_from_bytes_grow() -> Result<(), JsonError> { let s = &[0x00, 0x01, 0x02, 0x03]; - let mut jb = JsonBuilder::new_array(); + let mut jb = JsonBuilder::try_new_array().unwrap(); jb.append_string_from_bytes(s)?; for i in 1..1000 { let mut s = Vec::new(); - for _ in 0..i { - s.push(0x41); - } - let mut jb = JsonBuilder::new_array(); + s.resize(i, 0x41); + let mut jb = JsonBuilder::try_new_array().unwrap(); jb.append_string_from_bytes(&s)?; } @@ -1134,19 +1231,19 @@ mod test { #[test] fn test_invalid_utf8() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_string_from_bytes("invalid", &[0xf0, 0xf1, 0xf2]) .unwrap(); assert_eq!(jb.buf, r#"{"invalid":"\\xf0\\xf1\\xf2""#); - let mut jb = JsonBuilder::new_array(); + let mut jb = JsonBuilder::try_new_array().unwrap(); jb.append_string_from_bytes(&[0xf0, 0xf1, 0xf2]).unwrap(); assert_eq!(jb.buf, r#"["\\xf0\\xf1\\xf2""#); } #[test] fn test_marks() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_string("foo", "bar").unwrap(); assert_eq!(jb.buf, r#"{"foo":"bar""#); assert_eq!(jb.current_state(), State::ObjectNth); @@ -1169,7 +1266,7 @@ mod test { #[test] fn test_set_formatted() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_formatted("\"foo\":\"bar\"").unwrap(); assert_eq!(jb.buf, r#"{"foo":"bar""#); jb.set_formatted("\"bar\":\"foo\"").unwrap(); @@ -1180,7 +1277,7 @@ mod test { #[test] fn test_set_float() { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object().unwrap(); jb.set_float("one", 1.1).unwrap(); jb.set_float("two", 2.2).unwrap(); jb.close().unwrap(); @@ -1189,7 +1286,7 @@ mod test { #[test] fn test_append_float() { - let mut jb = JsonBuilder::new_array(); + let mut jb = JsonBuilder::try_new_array().unwrap(); jb.append_float(1.1).unwrap(); jb.append_float(2.2).unwrap(); jb.close().unwrap(); diff --git a/rust/src/pgsql/logger.rs b/rust/src/pgsql/logger.rs index 61c13a1092b3..43b661805a6b 100644 --- a/rust/src/pgsql/logger.rs +++ b/rust/src/pgsql/logger.rs @@ -46,7 +46,7 @@ fn log_pgsql(tx: &PgsqlTransaction, flags: u32, js: &mut JsonBuilder) -> Result< } fn log_request(req: &PgsqlFEMessage, flags: u32) -> Result { - let mut js = JsonBuilder::new_object(); + let mut js = JsonBuilder::try_new_object()?; match req { PgsqlFEMessage::StartupMessage(StartupPacket { length: _, @@ -108,7 +108,7 @@ fn log_request(req: &PgsqlFEMessage, flags: u32) -> Result Result { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object()?; let mut array_open = false; for response in &tx.responses { if let PgsqlBEMessage::ParameterStatus(msg) = response { @@ -268,7 +268,7 @@ fn log_error_notice_field_types( } fn log_startup_parameters(params: &PgsqlStartupParameters) -> Result { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object()?; // User is a mandatory field in a pgsql message jb.set_string_from_bytes("user", ¶ms.user.value)?; if let Some(parameters) = ¶ms.optional_params { @@ -284,7 +284,7 @@ fn log_startup_parameters(params: &PgsqlStartupParameters) -> Result Result { - let mut jb = JsonBuilder::new_object(); + let mut jb = JsonBuilder::try_new_object()?; jb.set_string_from_bytes(param.name.to_str(), ¶m.value)?; jb.close()?; Ok(jb)