From 33805b555d059950a80f743c642b63bb751377da Mon Sep 17 00:00:00 2001 From: TCeason Date: Mon, 25 Nov 2024 11:34:09 +0800 Subject: [PATCH] fix(query): add_hours function may panic if the argument is too big --- src/query/expression/src/utils/date_helper.rs | 7 +- src/query/functions/src/scalars/datetime.rs | 158 +++++++++++++++--- .../functions/02_0012_function_datetimes.test | 16 ++ 3 files changed, 154 insertions(+), 27 deletions(-) diff --git a/src/query/expression/src/utils/date_helper.rs b/src/query/expression/src/utils/date_helper.rs index a52a7ea6f1364..462934c3fdbe4 100644 --- a/src/query/expression/src/utils/date_helper.rs +++ b/src/query/expression/src/utils/date_helper.rs @@ -293,7 +293,7 @@ impl EvalDaysImpl { } pub fn eval_timestamp(date: i64, delta: impl AsPrimitive) -> i64 { - let mut value = date.wrapping_add(delta.as_() * MICROSECS_PER_DAY); + let mut value = date.wrapping_add(delta.as_().wrapping_mul(MICROSECS_PER_DAY)); clamp_timestamp(&mut value); value } @@ -311,12 +311,13 @@ pub struct EvalTimesImpl; impl EvalTimesImpl { pub fn eval_date(date: i32, delta: impl AsPrimitive, factor: i64) -> i32 { clamp_date( - (date as i64 * MICROSECS_PER_DAY).wrapping_add(delta.as_() * factor * MICROS_PER_SEC), + (date as i64 * MICROSECS_PER_DAY) + .wrapping_add(delta.as_().wrapping_mul(factor * MICROS_PER_SEC)), ) } pub fn eval_timestamp(us: i64, delta: impl AsPrimitive, factor: i64) -> i64 { - let mut ts = us.wrapping_add(delta.as_() * factor * MICROS_PER_SEC); + let mut ts = us.wrapping_add(delta.as_().wrapping_mul(factor * MICROS_PER_SEC)); clamp_timestamp(&mut ts); ts } diff --git a/src/query/functions/src/scalars/datetime.rs b/src/query/functions/src/scalars/datetime.rs index 23c7af7b8d488..44c7ed9e4b285 100644 --- a/src/query/functions/src/scalars/datetime.rs +++ b/src/query/functions/src/scalars/datetime.rs @@ -872,44 +872,74 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::(|date, delta, builder, ctx| { - match EvalYearsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalYearsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + +} + + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + + } + }), ); registry.register_passthrough_nullable_2_arg::( concat!($op, "_years"), - |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( |ts, delta, builder, ctx| { - match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalYearsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); registry.register_passthrough_nullable_2_arg::( concat!($op, "_quarters"), - |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::(|date, delta, builder, ctx| { - match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta} * 3) { + let delta_i32: Result = (delta*3).try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + } + + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }), ); registry.register_passthrough_nullable_2_arg::( @@ -918,13 +948,23 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( |ts, delta, builder, ctx| { - match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta} * 3) { + let delta_i32: Result = (delta*3).try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + } + + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); @@ -934,13 +974,22 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::(|date, delta, builder, ctx| { - match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalMonthsImpl::eval_date(date, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }), ); registry.register_passthrough_nullable_2_arg::( @@ -949,13 +998,23 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( |ts, delta, builder, ctx| { - match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta}) { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + match EvalMonthsImpl::eval_timestamp(ts, ctx.func_ctx.jiff_tz.clone(), $signed_wrapper!{delta_i32}) { Ok(t) => builder.push(t), Err(e) => { ctx.set_error(builder.len(), e); builder.push(0); }, } + } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } + }, ), ); @@ -964,8 +1023,18 @@ macro_rules! impl_register_arith_functions { concat!($op, "_days"), |_, _, _| FunctionDomain::MayThrow, - vectorize_with_builder_2_arg::(|date, delta, builder, _| { - builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta})) + vectorize_with_builder_2_arg::(|date, delta, builder, ctx| { + + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta_i32})); + } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }), ); registry.register_passthrough_nullable_2_arg::( @@ -973,8 +1042,16 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( - |ts, delta, builder, _| { - builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta})) + |ts, delta, builder, ctx| { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { +builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta_i32})); } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); @@ -983,9 +1060,17 @@ macro_rules! impl_register_arith_functions { concat!($op, "_weeks"), |_, _, _| FunctionDomain::MayThrow, - vectorize_with_builder_2_arg::(|date, delta, builder, _| { + vectorize_with_builder_2_arg::(|date, delta, builder, ctx| { let delta = 7 * delta; - builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta})) + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { +builder.push(EvalDaysImpl::eval_date(date, $signed_wrapper!{delta_i32})) } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }), ); registry.register_passthrough_nullable_2_arg::( @@ -993,9 +1078,17 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( - |ts, delta, builder, _| { + |ts, delta, builder, ctx| { let delta = 7 * delta; - builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta})) + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + builder.push(EvalDaysImpl::eval_timestamp(ts, $signed_wrapper!{delta_i32})) } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); @@ -1005,13 +1098,21 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( - |ts, delta, builder, _| { - let val = (ts as i64) * 24 * 3600 * MICROS_PER_SEC; + |ts, delta, builder, ctx| { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + let val = (ts as i64) * 24 * 3600 * MICROS_PER_SEC; builder.push(EvalTimesImpl::eval_timestamp( val, - $signed_wrapper!{delta}, + $signed_wrapper!{delta_i32}, FACTOR_HOUR, - )); + )); } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); @@ -1020,12 +1121,21 @@ macro_rules! impl_register_arith_functions { |_, _, _| FunctionDomain::MayThrow, vectorize_with_builder_2_arg::( - |ts, delta, builder, _| { - builder.push(EvalTimesImpl::eval_timestamp( + |ts, delta, builder, ctx| { + let delta_i32: Result = delta.try_into(); + match delta_i32 { + Ok(delta_i32) => { + builder.push(EvalTimesImpl::eval_timestamp( ts, - $signed_wrapper!{delta}, + $signed_wrapper!{delta_i32}, FACTOR_HOUR, - )); + )); + } + Err(_) => { + ctx.set_error(builder.len(), format!("Numeric value '{}' is out of range", delta)); + builder.push(0); + } + } }, ), ); diff --git a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test index c38571b09a419..5b88a15e82c12 100644 --- a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test +++ b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test @@ -1466,3 +1466,19 @@ query T select TRY_TO_TIMESTAMP(1, 0), TRY_TO_TIMESTAMP(1, null); ---- 1970-01-01 00:00:01.000000 NULL + +statement error 1006 +SELECT add_hours(to_date(710455), 1512263452497496403); + +statement error 1006 +SELECT add_hours(to_timestamp(710455), 1512263452497496403); + +query T +SELECT add_hours(to_date(710455), 2147483647); +---- +1000-01-01 00:00:00.000000 + +query T +SELECT add_hours(to_timestamp(710455), 2147483647); +---- +1000-01-01 00:00:00.000000