jsonbuilder: check buffer growth

Use try_reserve before growing the internal buffer, and the internal
state vector. This allows allocation errors to be caught and an error
returned instead of just aborting the process.

Ticket: #6057
pull/8941/head
Jason Ish 2 years ago committed by Victor Julien
parent 95cfc2b34f
commit 33827beae5

@ -399,7 +399,7 @@ pub fn dns_print_addr(addr: &Vec<u8>) -> std::string::String {
/// Log SOA section fields.
fn dns_log_soa(soa: &DNSRDataSOA) -> Result<JsonBuilder, JsonError> {
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<JsonBuilder, JsonError> {
/// Log SSHFP section fields.
fn dns_log_sshfp(sshfp: &DNSRDataSSHFP) -> Result<JsonBuilder, JsonError> {
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<JsonBuilder, JsonError> {
/// Log SRV section fields.
fn dns_log_srv(srv: &DNSRDataSRV) -> Result<JsonBuilder, JsonError> {
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<JsonBuilder, JsonError> {
}
fn dns_log_json_answer_detail(answer: &DNSAnswerEntry) -> Result<JsonBuilder, JsonError> {
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)?)?;

@ -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<TryReserveError> for JsonError {
fn from(_: TryReserveError) -> Self {
JsonError::Memory
}
}
impl From<Utf8Error> 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, JsonError> {
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<Self, JsonError> {
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, JsonError> {
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<Self, JsonError> {
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<String, JsonError> {
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();

@ -46,7 +46,7 @@ fn log_pgsql(tx: &PgsqlTransaction, flags: u32, js: &mut JsonBuilder) -> Result<
}
fn log_request(req: &PgsqlFEMessage, flags: u32) -> Result<JsonBuilder, JsonError> {
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<JsonBuilder, JsonErro
}
fn log_response_object(tx: &PgsqlTransaction) -> Result<JsonBuilder, JsonError> {
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<JsonBuilder, JsonError> {
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", &params.user.value)?;
if let Some(parameters) = &params.optional_params {
@ -284,7 +284,7 @@ fn log_startup_parameters(params: &PgsqlStartupParameters) -> Result<JsonBuilder
}
fn log_pgsql_param(param: &PgsqlParameter) -> Result<JsonBuilder, JsonError> {
let mut jb = JsonBuilder::new_object();
let mut jb = JsonBuilder::try_new_object()?;
jb.set_string_from_bytes(param.name.to_str(), &param.value)?;
jb.close()?;
Ok(jb)

Loading…
Cancel
Save