Skip to content

Commit e8f7e80

Browse files
mergify[bot]stdpain
mergify[bot]
and
stdpain
authored
[BugFix] Fix UAF when spill yield (backport #57311) (#57536)
Signed-off-by: stdpain <34912776+stdpain@users.noreply.github.com> Co-authored-by: stdpain <34912776+stdpain@users.noreply.github.com>
1 parent 0800c2d commit e8f7e80

6 files changed

Lines changed: 133 additions & 2 deletions

File tree

‎be/src/exec/pipeline/pipeline_driver_executor.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929

3030
namespace starrocks::pipeline {
3131

32+
DEFINE_FAIL_POINT(operator_return_failed_status);
33+
3234
GlobalDriverExecutor::GlobalDriverExecutor(const std::string& name, std::unique_ptr<ThreadPool> thread_pool,
3335
bool enable_resource_group, const CpuUtil::CpuIds& cpuids)
3436
: Base("pip_exec_" + name),
@@ -168,6 +170,12 @@ void GlobalDriverExecutor::_worker_thread() {
168170
status = driver->workgroup()->check_big_query(*query_ctx);
169171
}
170172

173+
FAIL_POINT_TRIGGER_EXECUTE(operator_return_failed_status, {
174+
if (status.ok()) {
175+
status = Status::InternalError("injected failed status");
176+
}
177+
});
178+
171179
if (!status.ok()) {
172180
auto o_id = get_backend_id();
173181
int64_t be_id = o_id.has_value() ? o_id.value() : -1;
@@ -177,6 +185,7 @@ void GlobalDriverExecutor::_worker_thread() {
177185
<< ", status=" << status;
178186
driver->runtime_profile()->add_info_string("ErrorMsg", std::string(status.message()));
179187
query_ctx->cancel(status);
188+
runtime_state->set_is_cancelled(true);
180189
driver->cancel_operators(runtime_state);
181190
if (driver->is_still_pending_finish()) {
182191
driver->set_driver_state(DriverState::PENDING_FINISH);

‎be/src/exec/spill/spill_components.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#pragma once
1616

1717
#include <functional>
18+
#include <memory>
1819
#include <queue>
1920

2021
#include "column/vectorized_fwd.h"
@@ -178,6 +179,8 @@ class RawSpillerWriter final : public SpillerWriter {
178179

179180
public:
180181
struct FlushContext : public SpillIOTaskContext {
182+
FlushContext(std::shared_ptr<Spiller> spiller_) : spiller(std::move(spiller_)) {}
183+
std::shared_ptr<Spiller> spiller;
181184
std::shared_ptr<SpillOutputDataStream> output;
182185
std::shared_ptr<BlockGroup> block_group;
183186
InputStreamPtr input_stream;
@@ -303,6 +306,7 @@ class PartitionedSpillerWriter final : public SpillerWriter {
303306

304307
public:
305308
struct PartitionedFlushContext : public SpillIOTaskContext {
309+
PartitionedFlushContext(std::shared_ptr<Spiller> spiller_) : spiller(std::move(spiller_)) {}
306310
// used in spill stage
307311
struct SpillStageContext {
308312
size_t processing_idx{};
@@ -327,6 +331,7 @@ class PartitionedSpillerWriter final : public SpillerWriter {
327331
PartitionedFlushContext(PartitionedFlushContext&&) = default;
328332
PartitionedFlushContext& operator=(PartitionedFlushContext&&) = default;
329333

334+
std::shared_ptr<Spiller> spiller;
330335
SpillStageContext spill_stage_ctx;
331336
SplitStageContext split_stage_ctx;
332337
};

‎be/src/exec/spill/spiller.hpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,8 @@ Status RawSpillerWriter::flush(RuntimeState* state, MemGuard&& guard) {
160160
DCHECK(has_pending_data());
161161
//
162162
if (!yield_ctx.task_context_data.has_value()) {
163-
yield_ctx.task_context_data = SpillIOTaskContextPtr(std::make_shared<FlushContext>());
163+
yield_ctx.task_context_data =
164+
SpillIOTaskContextPtr(std::make_shared<FlushContext>(_spiller->shared_from_this()));
164165
}
165166
auto defer = CancelableDefer([&]() {
166167
{
@@ -332,7 +333,8 @@ Status PartitionedSpillerWriter::flush(RuntimeState* state, bool is_final_flush,
332333
yield_ctx.time_spent_ns = 0;
333334
yield_ctx.need_yield = false;
334335
if (!yield_ctx.task_context_data.has_value()) {
335-
yield_ctx.task_context_data = SpillIOTaskContextPtr(std::make_shared<PartitionedFlushContext>());
336+
yield_ctx.task_context_data =
337+
SpillIOTaskContextPtr(std::make_shared<PartitionedFlushContext>(_spiller->shared_from_this()));
336338
}
337339
_spiller->update_spilled_task_status(
338340
yieldable_flush_task(yield_ctx, splitting_partitions, spilling_partitions));

‎be/test/io/spill_test.cpp‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -646,6 +646,79 @@ TEST_F(SpillTest, partition_process) {
646646
}
647647
}
648648

649+
struct PredoSyncExecutor {
650+
static std::function<void()> predo;
651+
static Status submit(workgroup::ScanTask task) {
652+
do {
653+
predo();
654+
task.run();
655+
} while (!task.is_finished());
656+
return Status::OK();
657+
}
658+
static void force_submit(workgroup::ScanTask task) { (void)submit(std::move(task)); }
659+
};
660+
std::function<void()> PredoSyncExecutor::predo;
661+
662+
TEST_F(SpillTest, partition_yield_with_failed) {
663+
ObjectPool pool;
664+
665+
// order by id_int
666+
// full data id_int, id_smallint
667+
std::vector<bool> nullables = {false, false};
668+
TExprBuilder tuple_slots_builder;
669+
tuple_slots_builder << TYPE_INT;
670+
auto tuple_slots = tuple_slots_builder.get_res();
671+
672+
auto ctx_st = no_partition_context(&pool, &dummy_rt_st, {}, tuple_slots);
673+
ASSERT_OK(ctx_st.status());
674+
auto ctx = ctx_st.value();
675+
(void)ctx;
676+
677+
std::vector<ExprContext*> tuple;
678+
ASSERT_OK(Expr::create_expr_trees(&pool, tuple_slots, &tuple, &dummy_rt_st));
679+
680+
// create chunk
681+
RandomChunkBuilder chunk_builder;
682+
683+
// create spilled factory
684+
// auto factory_options = SpilledFactoryOptions(ctx->partition_nums, ctx->parition_exprs, ctx->sort_exprs, ctx->sort_descs, false);
685+
auto factory = spill::make_spilled_factory();
686+
687+
// create spiller
688+
SpilledOptions spill_options(4);
689+
// 4 buffer chunk
690+
spill_options.mem_table_pool_size = 1;
691+
// file size: 1M
692+
spill_options.spill_mem_table_bytes_size = 1 * 1024 * 1024;
693+
// spill format type
694+
spill_options.spill_type = spill::SpillFormaterType::SPILL_BY_COLUMN;
695+
696+
spill_options.block_manager = dummy_block_mgr.get();
697+
698+
auto chunk_empty = chunk_builder.gen(tuple, nullables);
699+
700+
auto spiller = factory->create(spill_options);
701+
spiller->set_metrics(metrics);
702+
SpillerCaller<spill::PartitionedSpillerWriter*, spill::SpillerReader*> caller(spiller.get());
703+
ASSERT_OK(spiller->prepare(&dummy_rt_st));
704+
705+
size_t test_loop = 1024;
706+
std::vector<ChunkPtr> holder;
707+
{
708+
for (size_t i = 0; i < test_loop; ++i) {
709+
auto chunk = chunk_builder.gen(tuple, nullables);
710+
auto hash_column = spill::SpillHashColumn::create(chunk->num_rows());
711+
chunk->append_column(std::move(hash_column), -1);
712+
ASSERT_OK(spiller->spill<SyncExecutor>(&dummy_rt_st, chunk, EmptyMemGuard{}));
713+
ASSERT_OK(spiller->_spilled_task_status);
714+
holder.push_back(chunk);
715+
}
716+
717+
PredoSyncExecutor::predo = [&]() { spiller.reset(); };
718+
ASSERT_OK(spiller->flush<PredoSyncExecutor>(&dummy_rt_st, EmptyMemGuard{}));
719+
}
720+
}
721+
649722
TEST_F(SpillTest, aligned_buffer) {
650723
spill::AlignedBuffer buffer;
651724
ASSERT_EQ(buffer.data(), nullptr);

‎test/sql/test_exception/R/test_pipeline_operator_failed‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
-- name: test_pipeline_operator_failed @sequential
2+
create table t0 (
3+
c0 INT,
4+
c1 BIGINT
5+
) DUPLICATE KEY(c0) DISTRIBUTED BY HASH(c0) BUCKETS 1 PROPERTIES('replication_num' = '1');
6+
-- result:
7+
-- !result
8+
insert into t0 SELECT generate_series, 4096 - generate_series FROM TABLE(generate_series(1, 4096));
9+
-- result:
10+
-- !result
11+
set enable_spill=true;
12+
-- result:
13+
-- !result
14+
set spill_mode="force";
15+
-- result:
16+
-- !result
17+
admin enable failpoint 'operator_return_failed_status';
18+
-- result:
19+
-- !result
20+
[UC] select count(*) from t0;
21+
-- result:
22+
-- !result
23+
[UC] select count(*) from t0 group by c0;
24+
-- result:
25+
-- !result
26+
admin disable failpoint 'operator_return_failed_status';
27+
-- result:
28+
-- !result

‎test/sql/test_exception/T/test_pipeline_operator_failed‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
-- name: test_pipeline_operator_failed @sequential
2+
3+
create table t0 (
4+
c0 INT,
5+
c1 BIGINT
6+
) DUPLICATE KEY(c0) DISTRIBUTED BY HASH(c0) BUCKETS 1 PROPERTIES('replication_num' = '1');
7+
insert into t0 SELECT generate_series, 4096 - generate_series FROM TABLE(generate_series(1, 4096));
8+
9+
set enable_spill=true;
10+
set spill_mode="force";
11+
admin enable failpoint 'operator_return_failed_status';
12+
[UC] select count(*) from t0;
13+
[UC] select count(*) from t0 group by c0;
14+
admin disable failpoint 'operator_return_failed_status';

0 commit comments

Comments
 (0)