github-actions[bot] commented on code in PR #69257: URL: https://github.com/apache/doris/pull/69257#discussion_r4333722846
########## be/src/exprs/datasketches_hll_util.cpp: ########## @@ -0,0 +1,406 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "exprs/datasketches_hll_util.h" + +#include <bthread/bthread.h> + +#include <cstring> +#include <exception> +#include <mutex> +#include <new> +#include <stdexcept> + +#include "common/check.h" +#include "common/exception.h" +#include "common/status.h" +#include "util/debug_points.h" + +namespace doris::datasketches_hll { + +struct alignas(std::max_align_t) AllocationHeader { + AllocationHeader* previous; + AllocationHeader* next; + detail::AllocationOwner* owner; + size_t count; +}; + +namespace { + +using namespace datasketches::hll_constants; + +thread_local detail::AllocationOwner* pthread_scope = nullptr; + +bthread_key_t scope_key() { + static bthread_key_t key; + static std::once_flag once; + std::call_once(once, [] { DORIS_CHECK(bthread_key_create(&key, nullptr) == 0); }); + return key; +} + +detail::AllocationOwner* current_scope() { + return bthread_self() == 0 + ? pthread_scope + : static_cast<detail::AllocationOwner*>(bthread_getspecific(scope_key())); +} + +void set_scope(detail::AllocationOwner* scope) { + if (bthread_self() == 0) { + pthread_scope = scope; + } else { + DORIS_CHECK(bthread_setspecific(scope_key(), scope) == 0); + } +} + +class ActiveAllocationOwner { +public: + explicit ActiveAllocationOwner(detail::AllocationOwner* owner) : _previous(current_scope()) { + set_scope(owner); + } + ~ActiveAllocationOwner() { set_scope(_previous); } + +private: + detail::AllocationOwner* _previous; +}; + +// Reuse the library's protected static header readers without constructing a sketch or +// duplicating the mode/type bit decoding. This abstract adapter is never instantiated. +struct HeaderReader : datasketches::HllSketchImpl<Alloc> { + using datasketches::HllSketchImpl<Alloc>::extractCurMode; + using datasketches::HllSketchImpl<Alloc>::extractTgtHllType; +}; + +uint32_t read_uint32(const uint8_t* bytes) { + uint32_t value; + std::memcpy(&value, bytes, sizeof(value)); + return value; +} + +void check_set_size(const uint8_t* bytes, size_t size, uint8_t lg_k, bool compact) { + datasketches::HllUtil<Alloc>::checkMemSize(HASH_SET_INT_ARR_START, size); + const auto count = read_uint32(bytes + HASH_SET_COUNT_INT); + // SET promotes to HLL above 3/4 of its largest table (K/8 entries). Bound the count before + // computeLgArrInts(), whose power-of-two and load-factor arithmetic uses uint32_t. + if (lg_k <= 7 || count > (RESIZE_NUMER * (uint32_t {1} << (lg_k - 3))) / RESIZE_DENOM) { + throw std::invalid_argument("Invalid HLL SET coupon count or lgK"); + } + size_t items = count; + if (!compact) { + // The library evaluates 1 << lgArr before checking the payload. For example, lgArr=32 + // is an undefined shift. Preserve its support for older images with an unset lgArr. + const uint8_t lg_arr = bytes[LG_ARR_BYTE] < LG_INIT_SET_SIZE + ? datasketches::HllUtil<Alloc>::computeLgArrInts( + datasketches::SET, count, lg_k) + : bytes[LG_ARR_BYTE]; + if (lg_arr > lg_k - 3) { + throw std::invalid_argument("Invalid HLL SET array exponent"); + } + items = size_t {1} << lg_arr; + } + datasketches::HllUtil<Alloc>::checkMemSize(HASH_SET_INT_ARR_START + items * sizeof(uint32_t), + size); Review Comment: [P1] Reject SET coupon counts that disagree with the payload A complete updatable SET image can pass this check while its declared coupon count is smaller than its occupied slots. For example, take an lgK=8 HLL_4 SET with ten coupons, change only `HASH_SET_COUNT_INT` to 1, and call `datasketches_hll_downsample(sk,8,'HLL_6')`. The pinned `CouponHashSet::newSet` stores count=1 but copies all 32 slots; the equal-lgK union and `get_result` preserve that mismatch. `CouponList::serialize` then allocates only `12 + 4*1 = 16` bytes for the compact result but iterates all ten nonzero coupons, writing 52 bytes and corrupting BE heap memory. Compare the actual occupied-slot count with the declared count before deserializing an updatable SET, and add a complete-payload negative case. This is separate from the existing auxiliary-map and exponent findings: the exponent and payload size here are valid. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
