diff --git a/rust/src/dns/dns.rs b/rust/src/dns/dns.rs index 67eb3085547f..264129359652 100644 --- a/rust/src/dns/dns.rs +++ b/rust/src/dns/dns.rs @@ -19,6 +19,7 @@ use std; use std::collections::HashMap; use std::collections::VecDeque; use std::ffi::CString; +use std::panic::catch_unwind; use crate::applayer::*; use crate::core::{self, *}; @@ -769,8 +770,10 @@ pub unsafe extern "C" fn rs_dns_parse_request( flow: *const core::Flow, state: *mut std::os::raw::c_void, _pstate: *mut std::os::raw::c_void, stream_slice: StreamSlice, _data: *const std::os::raw::c_void, ) -> AppLayerResult { - let state = cast_pointer!(state, DNSState); - state.parse_request_udp(flow, stream_slice); + let _ = catch_unwind(|| { + let state = cast_pointer!(state, DNSState); + state.parse_request_udp(flow, stream_slice); + }); AppLayerResult::ok() } @@ -779,8 +782,10 @@ pub unsafe extern "C" fn rs_dns_parse_response( flow: *const core::Flow, state: *mut std::os::raw::c_void, _pstate: *mut std::os::raw::c_void, stream_slice: StreamSlice, _data: *const std::os::raw::c_void, ) -> AppLayerResult { - let state = cast_pointer!(state, DNSState); - state.parse_response_udp(flow, stream_slice); + let _ = catch_unwind(|| { + let state = cast_pointer!(state, DNSState); + state.parse_response_udp(flow, stream_slice); + }); AppLayerResult::ok() } @@ -790,13 +795,18 @@ pub unsafe extern "C" fn rs_dns_parse_request_tcp( flow: *const core::Flow, state: *mut std::os::raw::c_void, _pstate: *mut std::os::raw::c_void, stream_slice: StreamSlice, _data: *const std::os::raw::c_void, ) -> AppLayerResult { - let state = cast_pointer!(state, DNSState); - if stream_slice.is_gap() { - state.request_gap(stream_slice.gap_size()); - } else if !stream_slice.is_empty() { - return state.parse_request_tcp(flow, stream_slice); + match catch_unwind(|| { + let state = cast_pointer!(state, DNSState); + if stream_slice.is_gap() { + state.request_gap(stream_slice.gap_size()); + } else if !stream_slice.is_empty() { + return state.parse_request_tcp(flow, stream_slice); + } + return AppLayerResult::ok(); + }) { + Ok(result) => result, + Result::Err(_) => AppLayerResult::err(), } - AppLayerResult::ok() } #[no_mangle] @@ -804,13 +814,21 @@ pub unsafe extern "C" fn rs_dns_parse_response_tcp( flow: *const core::Flow, state: *mut std::os::raw::c_void, _pstate: *mut std::os::raw::c_void, stream_slice: StreamSlice, _data: *const std::os::raw::c_void, ) -> AppLayerResult { - let state = cast_pointer!(state, DNSState); - if stream_slice.is_gap() { - state.response_gap(stream_slice.gap_size()); - } else if !stream_slice.is_empty() { - return state.parse_response_tcp(flow, stream_slice); + match catch_unwind(|| { + let state = cast_pointer!(state, DNSState); + if stream_slice.is_gap() { + state.response_gap(stream_slice.gap_size()); + } else if !stream_slice.is_empty() { + return state.parse_response_tcp(flow, stream_slice); + } + return AppLayerResult::ok(); + }) { + Ok(result) => result, + Result::Err(_) => { + println!("caught error!"); + return AppLayerResult::err(); + } } - AppLayerResult::ok() } #[no_mangle] @@ -938,24 +956,29 @@ pub unsafe extern "C" fn rs_dns_probe( pub unsafe extern "C" fn rs_dns_probe_tcp( _flow: *const core::Flow, direction: u8, input: *const u8, len: u32, rdir: *mut u8, ) -> AppProto { - if len == 0 || len < std::mem::size_of::() as u32 + 2 { - return core::ALPROTO_UNKNOWN; - } - let slice: &[u8] = std::slice::from_raw_parts(input as *mut u8, len as usize); - //is_incomplete is checked by caller - let (is_dns, is_request, _) = probe_tcp(slice); - if is_dns { - let dir = if is_request { - Direction::ToServer - } else { - Direction::ToClient - }; - if (direction & DIR_BOTH) != dir.into() { - *rdir = dir as u8; + match catch_unwind(|| { + if len == 0 || len < std::mem::size_of::() as u32 + 2 { + return core::ALPROTO_UNKNOWN; } - return ALPROTO_DNS; + let slice: &[u8] = std::slice::from_raw_parts(input as *mut u8, len as usize); + //is_incomplete is checked by caller + let (is_dns, is_request, _) = probe_tcp(slice); + if is_dns { + let dir = if is_request { + Direction::ToServer + } else { + Direction::ToClient + }; + if (direction & DIR_BOTH) != dir.into() { + *rdir = dir as u8; + } + return ALPROTO_DNS; + } + return 0; + }) { + Ok(proto) => proto, + Result::Err(_) => ALPROTO_UNKNOWN, } - return 0; } #[no_mangle] 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 721e90844a97..6676b52294b5 100644 --- a/rust/src/jsonbuilder.rs +++ b/rust/src/jsonbuilder.rs @@ -1,4 +1,4 @@ -/* Copyright (C) 2020 Open Information Security Foundation +/* Copyright (C) 2020-2023 Open Information Security Foundation * * You can copy, redistribute or modify this Program under the terms of * the GNU General Public License version 2 as published by the Free @@ -17,6 +17,7 @@ #![allow(clippy::missing_safety_doc)] +use std::collections::TryReserveError; use std::ffi::CStr; use std::os::raw::c_char; use std::str::Utf8Error; @@ -27,6 +28,7 @@ const INIT_SIZE: usize = 4096; pub enum JsonError { InvalidState, Utf8Error(Utf8Error), + Memory, } impl std::error::Error for JsonError {} @@ -36,10 +38,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,33 +107,35 @@ 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 { + Ok(Self { buf, state: vec![State::None, State::ObjectFirst], 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 { + Ok(Self { buf, state: vec![State::None, State::ArrayFirst], init_type: Type::Array, - } + }) } // Reset the builder to its initial state, without losing @@ -159,7 +170,7 @@ impl JsonBuilder { State::None => { debug_validate_fail!("invalid state"); Err(JsonError::InvalidState) - }, + } } } @@ -334,7 +345,7 @@ impl JsonBuilder { State::ArrayFirst => { self.buf.push('"'); for i in 0..val.len() { - self.buf.push(HEX[(val[i] >> 4) as usize] as char); + self.buf.push(HEX[(val[i] >> 4) as usize] as char); self.buf.push(HEX[(val[i] & 0xf) as usize] as char); } self.buf.push('"'); @@ -345,7 +356,7 @@ impl JsonBuilder { self.buf.push(','); self.buf.push('"'); for i in 0..val.len() { - self.buf.push(HEX[(val[i] >> 4) as usize] as char); + self.buf.push(HEX[(val[i] >> 4) as usize] as char); self.buf.push(HEX[(val[i] & 0xf) as usize] as char); } self.buf.push('"'); @@ -522,7 +533,7 @@ impl JsonBuilder { self.buf.push_str(key); self.buf.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] >> 4) as usize] as char); self.buf.push(HEX[(val[i] & 0xf) as usize] as char); } self.buf.push('"'); @@ -675,14 +686,24 @@ fn string_from_bytes(input: &[u8]) -> String { #[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 +929,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 +958,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 +981,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 +999,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 +1035,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 +1056,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 +1075,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 +1104,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 +1134,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 +1149,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 +1159,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 +1174,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 +1209,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 +1220,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 +1229,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)