Fix int overflow in Arrow text body buffer computation (#9807)

The PG text values can be up to 1GB in size, so the computation for
extending the Arrow text body buffer could overflow. Switch it to use
the Size type.

Author: akuzm

tsl/src/nodes/vector_agg/exec.c

@@ -72,7 +72,7 @@ typedef struct
 
 	uint64 *restrict validity;
 
-	int allocated_body_bytes;
+	Size allocated_body_bytes;
 	uint8 *restrict body_buffer;
 
 	uint32 *restrict offset_buffer;
@@ -171,9 +171,19 @@ columnar_result_set_row(ColumnarResult *columnar_result, DecompressBatchState co
 		}
 		case DT_ArrowText:
 		{
-			const int result_bytes = VARSIZE_ANY_EXHDR(datum);
-			const int required_body_bytes =
+			const Size result_bytes = VARSIZE_ANY_EXHDR(datum);
+			const Size required_body_bytes =
 				pad_to_multiple(64, columnar_result->current_offset + result_bytes);
+			if (required_body_bytes > MaxAllocSize)
+			{
+				/*
+				 * The body buffer is allocated through palloc, so it cannot
+				 * exceed MaxAllocSize.
+				 */
+				ereport(ERROR,
+						(errcode(ERRCODE_PROGRAM_LIMIT_EXCEEDED),
+						 errmsg("vectorized text result exceeds %zu bytes", (Size) MaxAllocSize)));
+			}
 			if (required_body_bytes > columnar_result->allocated_body_bytes)
 			{
 				/*
@@ -183,8 +193,9 @@ columnar_result_set_row(ColumnarResult *columnar_result, DecompressBatchState co
 				 * tuned manually on a few real data sets until this balance
 				 * looked somewhat acceptable.
 				 */
-				const int new_body_bytes =
-					required_body_bytes * Min(10, Max(1.2, 1.2 * nrows / ((float) row + 1))) + 1;
+				const double desired_growth_factor = 1.2 * nrows / ((double) row + 1);
+				const double growth_factor = Min(10, Max(1.2, desired_growth_factor));
+				Size new_body_bytes = Min(MaxAllocSize, (Size) required_body_bytes * growth_factor);
 				Assert(new_body_bytes >= required_body_bytes);
 				columnar_result->body_buffer =
 					repalloc(columnar_result->body_buffer, new_body_bytes);

tsl/test/expected/vector_agg_text_overflow.out

@@ -0,0 +1,60 @@
+-- This file and its contents are licensed under the Timescale License.
+-- Please see the included NOTICE for copyright information and
+-- LICENSE-TIMESCALE for a copy of the license.
+-- Tests for body buffer overflow in columnar_result_set_row with
+-- DT_ArrowText type. When vectorized text functions produce large
+-- results, the body buffer growth computation could overflow int32,
+-- and the total buffer size can exceed MaxAllocSize.
+set max_parallel_workers_per_gather to 0;
+create table t_textoverflow (time timestamptz, a text);
+select table_name from create_hypertable('t_textoverflow', 'time',
+    chunk_time_interval => interval '1 year');
+   table_name   
+----------------
+ t_textoverflow
+
+insert into t_textoverflow
+select '2020-01-01'::timestamptz + i * interval '1 hour',
+       repeat('x', 100000)
+from generate_series(1, 10) i;
+alter table t_textoverflow set (
+    timescaledb.compress,
+    timescaledb.compress_orderby = 'time'
+);
+select count(compress_chunk(c)) from show_chunks('t_textoverflow') c;
+ count 
+-------
+     1
+
+set timescaledb.debug_require_vector_agg to 'require';
+select sum(length(a || a)) from t_textoverflow;
+   sum   
+---------
+ 2000000
+
+reset timescaledb.debug_require_vector_agg;
+drop table t_textoverflow cascade;
+-- When the actual data doesn't fit into MaxAllocSize, produce an error message
+-- instead of erroring out in palloc.
+create table t_text_overflow (time timestamptz not null, text_col text not null);
+select table_name from create_hypertable('t_text_overflow', 'time',
+    chunk_time_interval => interval '1 year');
+   table_name    
+-----------------
+ t_text_overflow
+
+insert into t_text_overflow
+select '2020-01-01'::timestamptz + i * interval '1 hour', 'a'
+from generate_series(1, 1000) i;
+alter table t_text_overflow set (
+    timescaledb.compress,
+    timescaledb.compress_orderby = 'time'
+);
+select count(compress_chunk(c)) from show_chunks('t_text_overflow') c;
+ count 
+-------
+     1
+
+set timescaledb.debug_require_vector_agg to 'require';
+select count(text_col || repeat('x', 2000000)) from t_text_overflow;
+ERROR:  vectorized text result exceeds 1073741823 bytes

tsl/test/sql/CMakeLists.txt

@@ -193,6 +193,7 @@ if(CMAKE_BUILD_TYPE MATCHES Debug)
     vector_agg_filter.sql
     vector_agg_grouping.sql
     vector_agg_text.sql
+    vector_agg_text_overflow.sql
     vector_agg_memory.sql
     vector_agg_modify_hypertable.sql
     vector_agg_segmentby.sql

tsl/test/sql/vector_agg_text_overflow.sql

@@ -0,0 +1,57 @@
+-- This file and its contents are licensed under the Timescale License.
+-- Please see the included NOTICE for copyright information and
+-- LICENSE-TIMESCALE for a copy of the license.
+
+-- Tests for body buffer overflow in columnar_result_set_row with
+-- DT_ArrowText type. When vectorized text functions produce large
+-- results, the body buffer growth computation could overflow int32,
+-- and the total buffer size can exceed MaxAllocSize.
+
+set max_parallel_workers_per_gather to 0;
+
+create table t_textoverflow (time timestamptz, a text);
+
+select table_name from create_hypertable('t_textoverflow', 'time',
+    chunk_time_interval => interval '1 year');
+
+insert into t_textoverflow
+select '2020-01-01'::timestamptz + i * interval '1 hour',
+       repeat('x', 100000)
+from generate_series(1, 10) i;
+
+alter table t_textoverflow set (
+    timescaledb.compress,
+    timescaledb.compress_orderby = 'time'
+);
+select count(compress_chunk(c)) from show_chunks('t_textoverflow') c;
+
+set timescaledb.debug_require_vector_agg to 'require';
+select sum(length(a || a)) from t_textoverflow;
+reset timescaledb.debug_require_vector_agg;
+
+drop table t_textoverflow cascade;
+
+-- When the actual data doesn't fit into MaxAllocSize, produce an error message
+-- instead of erroring out in palloc.
+create table t_text_overflow (time timestamptz not null, text_col text not null);
+
+select table_name from create_hypertable('t_text_overflow', 'time',
+    chunk_time_interval => interval '1 year');
+
+insert into t_text_overflow
+select '2020-01-01'::timestamptz + i * interval '1 hour', 'a'
+from generate_series(1, 1000) i;
+
+alter table t_text_overflow set (
+    timescaledb.compress,
+    timescaledb.compress_orderby = 'time'
+);
+select count(compress_chunk(c)) from show_chunks('t_text_overflow') c;
+
+set timescaledb.debug_require_vector_agg to 'require';
+select count(text_col || repeat('x', 2000000)) from t_text_overflow;
+reset timescaledb.debug_require_vector_agg;
+
+drop table t_text_overflow cascade;
+
+reset max_parallel_workers_per_gather;