[fix](vec) fix block mem use-after-free bug in agg table read (apache…

…#7944)
amorynan · Feb 5, 2022 · 9eb1d1d · 9eb1d1d
1 parent 51abaa8
commit 9eb1d1d
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 18 deletions.
diff --git a/be/src/vec/olap/block_reader.cpp b/be/src/vec/olap/block_reader.cpp
@@ -266,7 +266,7 @@ OLAPStatus BlockReader::_unique_key_next_block(Block* block, MemPool* mem_pool,
 }
 
 void BlockReader::_insert_data_normal(MutableColumns& columns) {
-    auto block = _next_row.block;
+    auto block = _next_row.block.get();
     for (auto idx : _normal_columns_idx) {
         columns[_return_columns_loc[idx]]->insert_from(*block->get_by_position(idx).column,
                                                        _next_row.row_pos);
@@ -314,7 +314,7 @@ size_t BlockReader::_copy_agg_data() {
 
     for (size_t i = 0; i < copy_size; i++) {
         auto& ref = _stored_row_ref[i];
-        _temp_ref_map[ref.block].emplace_back(ref.row_pos, i);
+        _temp_ref_map[ref.block.get()].emplace_back(ref.row_pos, i);
     }
 
     for (auto idx : _agg_columns_idx) {
@@ -328,9 +328,11 @@ size_t BlockReader::_copy_agg_data() {
             }
         } else {
             for (auto& it : _temp_ref_map) {
-                auto& src_column = *it.first->get_by_position(idx).column;
-                for (auto& pos : it.second) {
-                    dst_column->replace_column_data(src_column, pos.first, pos.second);
+                if (!it.second.empty()) {
+                    auto& src_column = *it.first->get_by_position(idx).column;
+                    for (auto &pos : it.second) {
+                        dst_column->replace_column_data(src_column, pos.first, pos.second);
+                    }
                 }
             }
         }

diff --git a/be/src/vec/olap/block_reader.h b/be/src/vec/olap/block_reader.h
@@ -86,7 +86,7 @@ class BlockReader final : public TabletReader {
     void _update_agg_value(MutableColumns& columns, int begin, int end, bool is_close = true);
 
     VCollectIterator _vcollect_iter;
-    IteratorRowRef _next_row {nullptr, -1, false};
+    IteratorRowRef _next_row {{}, -1, false};
 
     std::vector<AggregateFunctionPtr> _agg_functions;
     std::vector<AggregateDataPtr> _agg_places;

diff --git a/be/src/vec/olap/vcollect_iterator.cpp b/be/src/vec/olap/vcollect_iterator.cpp
@@ -16,6 +16,7 @@
 // under the License.
 
 #include "vec/olap/vcollect_iterator.h"
+#include <memory>
 
 #include "olap/rowset/beta_rowset_reader.h"
 
@@ -163,8 +164,8 @@ OLAPStatus VCollectIterator::next(Block* block) {
 VCollectIterator::Level0Iterator::Level0Iterator(RowsetReaderSharedPtr rs_reader, TabletReader* reader)
         : LevelIterator(reader), _rs_reader(rs_reader), _reader(reader) {
     DCHECK_EQ(RowsetTypePB::BETA_ROWSET, rs_reader->type());
-    _block = _schema.create_block(_reader->_return_columns);
-    _ref.block = &_block;
+    _block = std::make_shared<Block>(_schema.create_block(_reader->_return_columns));
+    _ref.block = _block;
     _ref.row_pos = 0;
     _ref.is_same = false;
 }
@@ -179,18 +180,18 @@ int64_t VCollectIterator::Level0Iterator::version() const {
 
 OLAPStatus VCollectIterator::Level0Iterator::_refresh_current_row() {
     do {
-        if (_block.rows() != 0 && _ref.row_pos < _block.rows()) {
+        if (_block->rows() != 0 && _ref.row_pos < _block->rows()) {
             return OLAP_SUCCESS;
         } else {
             _ref.is_same = false;
             _ref.row_pos = 0;
-            _block.clear_column_data();
-            auto res = _rs_reader->next_block(&_block);
+            _block->clear_column_data();
+            auto res = _rs_reader->next_block(_block.get());
             if (res != OLAP_SUCCESS) {
                 return res;
             }
         }
-    } while (_block.rows() != 0);
+    } while (_block->rows() != 0);
     _ref.row_pos = -1;
     return OLAP_ERR_DATA_EOF;
 }
@@ -323,7 +324,8 @@ OLAPStatus VCollectIterator::Level1Iterator::_merge_next(IteratorRowRef* ref) {
         return _merge_next(ref);
     }
 
-    *ref = _ref = *_cur_child->current_row_ref();
+    _ref = *_cur_child->current_row_ref();
+    *ref = _ref;
 
     _cur_child->set_same(false);
 
@@ -341,9 +343,7 @@ OLAPStatus VCollectIterator::Level1Iterator::_normal_next(IteratorRowRef* ref) {
         _children.pop_front();
         if (!_children.empty()) {
             _cur_child = *(_children.begin());
-            auto result = _cur_child->next(ref);
-            _ref = *ref;
-            return result;
+            return _normal_next(ref);
         } else {
             _cur_child = nullptr;
             return OLAP_ERR_DATA_EOF;

diff --git a/be/src/vec/olap/vcollect_iterator.h b/be/src/vec/olap/vcollect_iterator.h
@@ -35,7 +35,7 @@ class TabletSchema;
 namespace vectorized {
 
 struct IteratorRowRef {
-    const Block* block;
+    std::shared_ptr<Block> block;
     int16_t row_pos;
     bool is_same;
 };
@@ -137,7 +137,7 @@ class VCollectIterator {
 
         RowsetReaderSharedPtr _rs_reader;
         TabletReader* _reader = nullptr;
-        Block _block;
+        std::shared_ptr<Block> _block;
     };
 
     // Iterate from LevelIterators (maybe Level0Iterators or Level1Iterator or mixed)