From 98b550e37f13d6009c0016ee493816271d3c4d90 Mon Sep 17 00:00:00 2001 From: baishen Date: Thu, 18 Jul 2024 16:53:29 +0800 Subject: [PATCH] fix: fix jsonpath selector unwrap panic --- benches/get_path.rs | 2 +- src/error.rs | 6 ++ src/functions.rs | 25 +++--- src/jsonpath/selector.rs | 173 +++++++++++++++++++++++---------------- tests/it/functions.rs | 10 ++- 5 files changed, 131 insertions(+), 85 deletions(-) diff --git a/benches/get_path.rs b/benches/get_path.rs index fbe3ff5..2211298 100644 --- a/benches/get_path.rs +++ b/benches/get_path.rs @@ -27,7 +27,7 @@ fn jsonb_get(data: &[u8], paths: &[&str], expected: &str) { let mut result_data = vec![]; let mut result_offsets = vec![]; - jsonb::get_by_path(data, json_path, &mut result_data, &mut result_offsets); + jsonb::get_by_path(data, json_path, &mut result_data, &mut result_offsets).unwrap(); let s = jsonb::as_str(&result_data).unwrap(); assert_eq!(s, expected); diff --git a/src/error.rs b/src/error.rs index db8a490..2dce541 100644 --- a/src/error.rs +++ b/src/error.rs @@ -109,3 +109,9 @@ impl From for Error { Error::InvalidUtf8 } } + +impl From>> for Error { + fn from(_error: nom::Err>) -> Self { + Error::InvalidJsonb + } +} diff --git a/src/functions.rs b/src/functions.rs index 986f0cb..032d182 100644 --- a/src/functions.rs +++ b/src/functions.rs @@ -155,7 +155,7 @@ pub fn array_length(value: &[u8]) -> Option { } /// Checks whether the JSON path returns any item for the `JSONB` value. -pub fn path_exists<'a>(value: &'a [u8], json_path: JsonPath<'a>) -> bool { +pub fn path_exists<'a>(value: &'a [u8], json_path: JsonPath<'a>) -> Result { let selector = Selector::new(json_path, Mode::Mixed); if !is_jsonb(value) { match parse_value(value) { @@ -163,7 +163,7 @@ pub fn path_exists<'a>(value: &'a [u8], json_path: JsonPath<'a>) -> bool { let value = val.to_vec(); selector.exists(value.as_slice()) } - Err(_) => false, + Err(_) => Ok(false), } } else { selector.exists(value) @@ -188,16 +188,17 @@ pub fn get_by_path<'a>( json_path: JsonPath<'a>, data: &mut Vec, offsets: &mut Vec, -) { +) -> Result<(), Error> { let selector = Selector::new(json_path, Mode::Mixed); if !is_jsonb(value) { if let Ok(val) = parse_value(value) { let value = val.to_vec(); - selector.select(value.as_slice(), data, offsets) + selector.select(value.as_slice(), data, offsets)?; } } else { - selector.select(value, data, offsets) + selector.select(value, data, offsets)?; } + Ok(()) } /// Get the inner element of `JSONB` value by JSON path. @@ -207,16 +208,17 @@ pub fn get_by_path_first<'a>( json_path: JsonPath<'a>, data: &mut Vec, offsets: &mut Vec, -) { +) -> Result<(), Error> { let selector = Selector::new(json_path, Mode::First); if !is_jsonb(value) { if let Ok(val) = parse_value(value) { let value = val.to_vec(); - selector.select(value.as_slice(), data, offsets) + selector.select(value.as_slice(), data, offsets)?; } } else { - selector.select(value, data, offsets) + selector.select(value, data, offsets)?; } + Ok(()) } /// Get the inner elements of `JSONB` value by JSON path. @@ -226,16 +228,17 @@ pub fn get_by_path_array<'a>( json_path: JsonPath<'a>, data: &mut Vec, offsets: &mut Vec, -) { +) -> Result<(), Error> { let selector = Selector::new(json_path, Mode::Array); if !is_jsonb(value) { if let Ok(val) = parse_value(value) { let value = val.to_vec(); - selector.select(value.as_slice(), data, offsets) + selector.select(value.as_slice(), data, offsets)?; } } else { - selector.select(value, data, offsets) + selector.select(value, data, offsets)?; } + Ok(()) } /// Get the inner element of `JSONB` Array by index. diff --git a/src/jsonpath/selector.rs b/src/jsonpath/selector.rs index 8988c86..19ea1bf 100644 --- a/src/jsonpath/selector.rs +++ b/src/jsonpath/selector.rs @@ -74,44 +74,50 @@ impl<'a> Selector<'a> { Self { json_path, mode } } - pub fn select(&'a self, root: &'a [u8], data: &mut Vec, offsets: &mut Vec) { - let mut poses = self.find_positions(root, None, &self.json_path.paths); + pub fn select( + &'a self, + root: &'a [u8], + data: &mut Vec, + offsets: &mut Vec, + ) -> Result<(), Error> { + let mut poses = self.find_positions(root, None, &self.json_path.paths)?; if self.json_path.is_predicate() { - Self::build_predicate_result(&mut poses, data); - return; + Self::build_predicate_result(&mut poses, data)?; + return Ok(()); } match self.mode { - Mode::All => Self::build_values(root, &mut poses, data, offsets), + Mode::All => Self::build_values(root, &mut poses, data, offsets)?, Mode::First => { poses.truncate(1); - Self::build_values(root, &mut poses, data, offsets) + Self::build_values(root, &mut poses, data, offsets)? } - Mode::Array => Self::build_scalar_array(root, &mut poses, data, offsets), + Mode::Array => Self::build_scalar_array(root, &mut poses, data, offsets)?, Mode::Mixed => { if poses.len() > 1 { - Self::build_scalar_array(root, &mut poses, data, offsets) + Self::build_scalar_array(root, &mut poses, data, offsets)? } else { - Self::build_values(root, &mut poses, data, offsets) + Self::build_values(root, &mut poses, data, offsets)? } } } + Ok(()) } - pub fn exists(&'a self, root: &'a [u8]) -> bool { + pub fn exists(&'a self, root: &'a [u8]) -> Result { if self.json_path.is_predicate() { - return true; + return Ok(true); } - let poses = self.find_positions(root, None, &self.json_path.paths); - !poses.is_empty() + let poses = self.find_positions(root, None, &self.json_path.paths)?; + Ok(!poses.is_empty()) } pub fn predicate_match(&'a self, root: &'a [u8]) -> Result { if !self.json_path.is_predicate() { return Err(Error::InvalidJsonPathPredicate); } - let poses = self.find_positions(root, None, &self.json_path.paths); + let poses = self.find_positions(root, None, &self.json_path.paths)?; Ok(!poses.is_empty()) } @@ -120,7 +126,7 @@ impl<'a> Selector<'a> { root: &'a [u8], current: Option<&Position>, paths: &[Path<'a>], - ) -> VecDeque { + ) -> Result, Error> { let mut poses = VecDeque::new(); let start_pos = if let Some(Path::Current) = paths.first() { @@ -139,7 +145,8 @@ impl<'a> Selector<'a> { let len = poses.len(); for _ in 0..len { let pos = poses.pop_front().unwrap(); - if self.filter_expr(root, &pos, expr) { + let res = self.filter_expr(root, &pos, expr)?; + if res { poses.push_back(pos); } } @@ -150,7 +157,7 @@ impl<'a> Selector<'a> { let pos = poses.pop_front().unwrap(); match pos { Position::Container((offset, length)) => { - self.select_path(root, offset, length, path, &mut poses); + self.select_path(root, offset, length, path, &mut poses)?; } Position::Scalar(_) => { // In lax mode, bracket wildcard allow Scalar value. @@ -163,7 +170,7 @@ impl<'a> Selector<'a> { } } } - poses + Ok(poses) } fn select_path( @@ -173,22 +180,23 @@ impl<'a> Selector<'a> { length: usize, path: &Path<'a>, poses: &mut VecDeque, - ) { + ) -> Result<(), Error> { match path { Path::DotWildcard => { - self.select_object_values(root, offset, poses); + self.select_object_values(root, offset, poses)?; } Path::BracketWildcard => { - self.select_array_values(root, offset, length, poses); + self.select_array_values(root, offset, length, poses)?; } Path::ColonField(name) | Path::DotField(name) | Path::ObjectField(name) => { - self.select_by_name(root, offset, name, poses); + self.select_by_name(root, offset, name, poses)?; } Path::ArrayIndices(indices) => { - self.select_by_indices(root, offset, indices, poses); + self.select_by_indices(root, offset, indices, poses)?; } _ => unreachable!(), } + Ok(()) } // select all values in an Object. @@ -197,13 +205,13 @@ impl<'a> Selector<'a> { root: &'a [u8], root_offset: usize, poses: &mut VecDeque, - ) { - let (rest, (ty, length)) = decode_header(&root[root_offset..]).unwrap(); + ) -> Result<(), Error> { + let (rest, (ty, length)) = decode_header(&root[root_offset..])?; if ty != OBJECT_CONTAINER_TAG || length == 0 { - return; + return Ok(()); } - let (rest, key_jentries) = decode_jentries(rest, length).unwrap(); - let (_, val_jentries) = decode_jentries(rest, length).unwrap(); + let (rest, key_jentries) = decode_jentries(rest, length)?; + let (_, val_jentries) = decode_jentries(rest, length)?; let mut offset = root_offset + 4 + length * 8; for (_, length) in key_jentries.iter() { offset += length; @@ -217,6 +225,7 @@ impl<'a> Selector<'a> { poses.push_back(pos); offset += jlength; } + Ok(()) } // select all values in an Array. @@ -226,14 +235,14 @@ impl<'a> Selector<'a> { root_offset: usize, root_length: usize, poses: &mut VecDeque, - ) { - let (rest, (ty, length)) = decode_header(&root[root_offset..]).unwrap(); + ) -> Result<(), Error> { + let (rest, (ty, length)) = decode_header(&root[root_offset..])?; if ty != ARRAY_CONTAINER_TAG { // In lax mode, bracket wildcard allow Scalar value. poses.push_back(Position::Container((root_offset, root_length))); - return; + return Ok(()); } - let (_, val_jentries) = decode_jentries(rest, length).unwrap(); + let (_, val_jentries) = decode_jentries(rest, length)?; let mut offset = root_offset + 4 + length * 4; for (jty, jlength) in val_jentries.iter() { let pos = if *jty == CONTAINER_TAG { @@ -244,6 +253,7 @@ impl<'a> Selector<'a> { poses.push_back(pos); offset += jlength; } + Ok(()) } // select value in an Object by key name. @@ -253,13 +263,13 @@ impl<'a> Selector<'a> { root_offset: usize, name: &str, poses: &mut VecDeque, - ) { - let (rest, (ty, length)) = decode_header(&root[root_offset..]).unwrap(); + ) -> Result<(), Error> { + let (rest, (ty, length)) = decode_header(&root[root_offset..])?; if ty != OBJECT_CONTAINER_TAG || length == 0 { - return; + return Ok(()); } - let (rest, key_jentries) = decode_jentries(rest, length).unwrap(); - let (_, val_jentries) = decode_jentries(rest, length).unwrap(); + let (rest, key_jentries) = decode_jentries(rest, length)?; + let (_, val_jentries) = decode_jentries(rest, length)?; let mut idx = 0; let mut offset = root_offset + 4 + length * 8; let mut found = false; @@ -268,7 +278,7 @@ impl<'a> Selector<'a> { offset += jlength; continue; } - let (_, key) = decode_string(&root[offset..], *jlength).unwrap(); + let (_, key) = decode_string(&root[offset..], *jlength)?; if name == unsafe { std::str::from_utf8_unchecked(key) } { found = true; idx = i; @@ -276,7 +286,7 @@ impl<'a> Selector<'a> { offset += jlength; } if !found { - return; + return Ok(()); } for (i, (jty, jlength)) in val_jentries.iter().enumerate() { if i != idx { @@ -291,6 +301,7 @@ impl<'a> Selector<'a> { poses.push_back(pos); break; } + Ok(()) } // select values in an Array by indices. @@ -300,10 +311,10 @@ impl<'a> Selector<'a> { root_offset: usize, indices: &Vec, poses: &mut VecDeque, - ) { - let (rest, (ty, length)) = decode_header(&root[root_offset..]).unwrap(); + ) -> Result<(), Error> { + let (rest, (ty, length)) = decode_header(&root[root_offset..])?; if ty != ARRAY_CONTAINER_TAG || length == 0 { - return; + return Ok(()); } let mut val_indices = Vec::new(); for index in indices { @@ -321,9 +332,9 @@ impl<'a> Selector<'a> { } } if val_indices.is_empty() { - return; + return Ok(()); } - let (_, jentries) = decode_jentries(rest, length).unwrap(); + let (_, jentries) = decode_jentries(rest, length)?; let mut offset = root_offset + 4 + length * 4; let mut offsets = Vec::with_capacity(jentries.len()); for (_, jlength) in jentries.iter() { @@ -340,15 +351,20 @@ impl<'a> Selector<'a> { }; poses.push_back(pos); } + Ok(()) } - fn build_predicate_result(poses: &mut VecDeque, data: &mut Vec) { + fn build_predicate_result( + poses: &mut VecDeque, + data: &mut Vec, + ) -> Result<(), Error> { let jentry = match poses.pop_front() { Some(_) => TRUE_TAG, None => FALSE_TAG, }; - data.write_u32::(SCALAR_CONTAINER_TAG).unwrap(); - data.write_u32::(jentry).unwrap(); + data.write_u32::(SCALAR_CONTAINER_TAG)?; + data.write_u32::(jentry)?; + Ok(()) } fn build_values( @@ -356,16 +372,16 @@ impl<'a> Selector<'a> { poses: &mut VecDeque, data: &mut Vec, offsets: &mut Vec, - ) { + ) -> Result<(), Error> { while let Some(pos) = poses.pop_front() { match pos { Position::Container((offset, length)) => { data.extend_from_slice(&root[offset..offset + length]); } Position::Scalar((ty, offset, length)) => { - data.write_u32::(SCALAR_CONTAINER_TAG).unwrap(); + data.write_u32::(SCALAR_CONTAINER_TAG)?; let jentry = ty | length as u32; - data.write_u32::(jentry).unwrap(); + data.write_u32::(jentry)?; if length > 0 { data.extend_from_slice(&root[offset..offset + length]); } @@ -373,6 +389,7 @@ impl<'a> Selector<'a> { } offsets.push(data.len() as u64); } + Ok(()) } fn build_scalar_array( @@ -380,11 +397,11 @@ impl<'a> Selector<'a> { poses: &mut VecDeque, data: &mut Vec, offsets: &mut Vec, - ) { + ) -> Result<(), Error> { let len = poses.len(); let header = ARRAY_CONTAINER_TAG | len as u32; // write header. - data.write_u32::(header).unwrap(); + data.write_u32::(header)?; let mut jentry_offset = data.len(); // reserve space for jentry. data.resize(jentry_offset + 4 * len, 0); @@ -407,6 +424,7 @@ impl<'a> Selector<'a> { jentry_offset += 4; } offsets.push(data.len() as u64); + Ok(()) } // check and convert index to Array index. @@ -445,23 +463,29 @@ impl<'a> Selector<'a> { } } - fn filter_expr(&'a self, root: &'a [u8], pos: &Position, expr: &Expr<'a>) -> bool { + fn filter_expr( + &'a self, + root: &'a [u8], + pos: &Position, + expr: &Expr<'a>, + ) -> Result { match expr { Expr::BinaryOp { op, left, right } => match op { BinaryOperator::Or => { - let lhs = self.filter_expr(root, pos, left); - let rhs = self.filter_expr(root, pos, right); - lhs || rhs + let lhs = self.filter_expr(root, pos, left)?; + let rhs = self.filter_expr(root, pos, right)?; + Ok(lhs || rhs) } BinaryOperator::And => { - let lhs = self.filter_expr(root, pos, left); - let rhs = self.filter_expr(root, pos, right); - lhs && rhs + let lhs = self.filter_expr(root, pos, left)?; + let rhs = self.filter_expr(root, pos, right)?; + Ok(lhs && rhs) } _ => { - let lhs = self.convert_expr_val(root, pos, *left.clone()); - let rhs = self.convert_expr_val(root, pos, *right.clone()); - self.compare(op, &lhs, &rhs) + let lhs = self.convert_expr_val(root, pos, *left.clone())?; + let rhs = self.convert_expr_val(root, pos, *right.clone())?; + let res = self.compare(op, &lhs, &rhs); + Ok(res) } }, Expr::FilterFunc(filter_expr) => match filter_expr { @@ -471,14 +495,25 @@ impl<'a> Selector<'a> { } } - fn eval_exists(&'a self, root: &'a [u8], pos: &Position, paths: &[Path<'a>]) -> bool { - let poses = self.find_positions(root, Some(pos), paths); - !poses.is_empty() + fn eval_exists( + &'a self, + root: &'a [u8], + pos: &Position, + paths: &[Path<'a>], + ) -> Result { + let poses = self.find_positions(root, Some(pos), paths)?; + let res = !poses.is_empty(); + Ok(res) } - fn convert_expr_val(&'a self, root: &'a [u8], pos: &Position, expr: Expr<'a>) -> ExprValue<'a> { + fn convert_expr_val( + &'a self, + root: &'a [u8], + pos: &Position, + expr: Expr<'a>, + ) -> Result, Error> { match expr { - Expr::Value(value) => ExprValue::Value(value.clone()), + Expr::Value(value) => Ok(ExprValue::Value(value.clone())), Expr::Paths(paths) => { // get value from path and convert to `ExprValue`. let mut poses = VecDeque::new(); @@ -500,7 +535,7 @@ impl<'a> Selector<'a> { let pos = poses.pop_front().unwrap(); match pos { Position::Container((offset, length)) => { - self.select_path(root, offset, length, path, &mut poses); + self.select_path(root, offset, length, path, &mut poses)?; } Position::Scalar(_) => { // In lax mode, bracket wildcard allow Scalar value. @@ -535,7 +570,7 @@ impl<'a> Selector<'a> { values.push(value); } } - ExprValue::Values(values) + Ok(ExprValue::Values(values)) } _ => unreachable!(), } diff --git a/tests/it/functions.rs b/tests/it/functions.rs index cc7c2bb..f273ff8 100644 --- a/tests/it/functions.rs +++ b/tests/it/functions.rs @@ -164,13 +164,13 @@ fn test_path_exists() { let value = parse_value(json.as_bytes()).unwrap().to_vec(); let json_path = parse_json_path(path.as_bytes()).unwrap(); let res = path_exists(value.as_slice(), json_path); - assert_eq!(res, expect); + assert_eq!(res, Ok(expect)); } // Check from String JSON { let json_path = parse_json_path(path.as_bytes()).unwrap(); let res = path_exists(json.as_bytes(), json_path); - assert_eq!(res, expect); + assert_eq!(res, Ok(expect)); } } } @@ -236,7 +236,8 @@ fn test_path_exists_expr() { let mut out_offsets: Vec = Vec::new(); let json_path = parse_json_path(path.as_bytes()).unwrap(); - get_by_path_array(&buf, json_path, &mut out_buf, &mut out_offsets); + let res = get_by_path_array(&buf, json_path, &mut out_buf, &mut out_offsets); + assert!(res.is_ok()); let expected_buf = parse_value(expected.as_bytes()).unwrap().to_vec(); assert_eq!(out_buf, expected_buf); @@ -311,7 +312,8 @@ fn test_get_by_path() { out_buf.clear(); out_offsets.clear(); let json_path = parse_json_path(path.as_bytes()).unwrap(); - get_by_path(&buf, json_path, &mut out_buf, &mut out_offsets); + let res = get_by_path(&buf, json_path, &mut out_buf, &mut out_offsets); + assert!(res.is_ok()); if expects.is_empty() { assert_eq!(out_offsets.len(), expects.len()); } else if expects.len() == 1 {