diff --git a/components/tidb_query/src/rpn_expr/impl_encryption.rs b/components/tidb_query/src/rpn_expr/impl_encryption.rs index 4a8b2ca02b6..85b181d7129 100644 --- a/components/tidb_query/src/rpn_expr/impl_encryption.rs +++ b/components/tidb_query/src/rpn_expr/impl_encryption.rs @@ -3,6 +3,8 @@ use openssl::hash::{self, MessageDigest}; use tidb_query_codegen::rpn_fn; +use super::super::expr::{Error, EvalContext}; + use crate::codec::data_type::*; use crate::Result; @@ -31,6 +33,22 @@ fn hex_digest(hashtype: MessageDigest, input: &[u8]) -> Result { .map_err(|e| box_err!("OpenSSL error: {:?}", e)) } +#[rpn_fn(capture = [ctx])] +#[inline] +pub fn uncompressed_length(ctx: &mut EvalContext, arg: &Option) -> Result> { + use byteorder::{ByteOrder, LittleEndian}; + Ok(arg.as_ref().map(|s| { + if s.is_empty() { + 0 + } else if s.len() <= 4 { + ctx.warnings.append_warning(Error::zlib_data_corrupted()); + 0 + } else { + Int::from(LittleEndian::read_u32(&s[0..4])) + } + })) +} + #[cfg(test)] mod tests { use tipb::ScalarFuncSig; @@ -53,6 +71,8 @@ mod tests { ); } + use hex; + #[test] fn test_md5() { let test_cases = vec![ @@ -117,4 +137,37 @@ mod tests { } test_unary_func_ok_none::(ScalarFuncSig::Sha1); } + + #[test] + fn test_uncompressed_length() { + let cases = vec![ + (Some(""), Some(0)), + ( + Some("0B000000789CCB48CDC9C95728CF2FCA4901001A0B045D"), + Some(11), + ), + ( + Some("0C000000789CCB48CDC9C95728CF2F32303402001D8004202E"), + Some(12), + ), + (Some("020000000000"), Some(2)), + (Some("0000000001"), Some(0)), + ( + Some("02000000789CCB48CDC9C95728CF2FCA4901001A0B045D"), + Some(2), + ), + (Some("010203"), Some(0)), + (Some("01020304"), Some(0)), + (None, None), + ]; + + for (s, exp) in cases { + let s = s.map(|inner| hex::decode(inner.as_bytes().to_vec()).unwrap()); + let output = RpnFnScalarEvaluator::new() + .push_param(s) + .evaluate(ScalarFuncSig::UncompressedLength) + .unwrap(); + assert_eq!(output, exp); + } + } } diff --git a/components/tidb_query/src/rpn_expr/mod.rs b/components/tidb_query/src/rpn_expr/mod.rs index 44919b48490..5ad9e8bfe0c 100644 --- a/components/tidb_query/src/rpn_expr/mod.rs +++ b/components/tidb_query/src/rpn_expr/mod.rs @@ -253,6 +253,7 @@ fn map_expr_node_to_rpn_func(expr: &Expr) -> Result { ScalarFuncSig::CaseWhenJson => case_when_fn_meta::(), ScalarFuncSig::Md5 => md5_fn_meta(), ScalarFuncSig::Sha1 => sha1_fn_meta(), + ScalarFuncSig::UncompressedLength => uncompressed_length_fn_meta(), ScalarFuncSig::DateFormatSig => date_format_fn_meta(), ScalarFuncSig::DayOfYear => day_of_year_fn_meta(), ScalarFuncSig::WeekDay => week_day_fn_meta(),