From ff6c3c75c2f4577f54dda7c47f6ad7db278e1784 Mon Sep 17 00:00:00 2001 From: Alexander Tokmakov Date: Tue, 12 Jan 2021 17:34:50 +0300 Subject: [PATCH 1/2] add protection from unsafe allocations --- base/daemon/BaseDaemon.cpp | 3 +++ src/Common/MemoryTracker.cpp | 12 ++++++++++++ src/Common/MemoryTracker.h | 18 ++++++++++++++++++ src/Common/QueryProfiler.cpp | 2 ++ src/Common/ThreadFuzzer.cpp | 1 + src/Common/ThreadPool.cpp | 4 +++- .../System/StorageSystemStackTrace.cpp | 1 + 7 files changed, 40 insertions(+), 1 deletion(-) diff --git a/base/daemon/BaseDaemon.cpp b/base/daemon/BaseDaemon.cpp index 830e7857a1f..4cf8a8d7ce9 100644 --- a/base/daemon/BaseDaemon.cpp +++ b/base/daemon/BaseDaemon.cpp @@ -112,11 +112,13 @@ static void writeSignalIDtoSignalPipe(int sig) /** Signal handler for HUP / USR1 */ static void closeLogsSignalHandler(int sig, siginfo_t *, void *) { + DENY_ALLOCATIONS_IN_SCOPE; writeSignalIDtoSignalPipe(sig); } static void terminateRequestedSignalHandler(int sig, siginfo_t *, void *) { + DENY_ALLOCATIONS_IN_SCOPE; writeSignalIDtoSignalPipe(sig); } @@ -125,6 +127,7 @@ static void terminateRequestedSignalHandler(int sig, siginfo_t *, void *) */ static void signalHandler(int sig, siginfo_t * info, void * context) { + DENY_ALLOCATIONS_IN_SCOPE; auto saved_errno = errno; /// We must restore previous value of errno in signal handler. char buf[signal_pipe_buf_size]; diff --git a/src/Common/MemoryTracker.cpp b/src/Common/MemoryTracker.cpp index 4b0e1f9cada..e62f15d4fd1 100644 --- a/src/Common/MemoryTracker.cpp +++ b/src/Common/MemoryTracker.cpp @@ -12,6 +12,10 @@ #include #include +#ifdef MEMORY_TRACKER_DEBUG_CHECKS +thread_local bool _memory_tracker_always_throw_logical_error_on_allocation = false; +#endif + namespace { @@ -165,6 +169,14 @@ void MemoryTracker::alloc(Int64 size) } } +#ifdef MEMORY_TRACKER_DEBUG_CHECKS + if (unlikely(_memory_tracker_always_throw_logical_error_on_allocation)) + { + _memory_tracker_always_throw_logical_error_on_allocation = false; + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Memory tracker: allocations not allowed."); + } +#endif + std::bernoulli_distribution fault(fault_probability); if (unlikely(fault_probability && fault(thread_local_rng)) && memoryTrackerCanThrow(level, true)) { diff --git a/src/Common/MemoryTracker.h b/src/Common/MemoryTracker.h index bfe03d19a27..b67f9e368e2 100644 --- a/src/Common/MemoryTracker.h +++ b/src/Common/MemoryTracker.h @@ -5,6 +5,24 @@ #include #include +#if !defined(NDEBUG) || defined(ADDRESS_SANITIZER) || defined(THREAD_SANITIZER) || defined(MEMORY_SANITIZER) || defined(UNDEFINED_BEHAVIOR_SANITIZER) +#define MEMORY_TRACKER_DEBUG_CHECKS +#endif + +#ifdef MEMORY_TRACKER_DEBUG_CHECKS +#include +extern thread_local bool _memory_tracker_always_throw_logical_error_on_allocation; +#define ALLOCATIONS_IN_SCOPE_IMPL_CONCAT(n, val) \ + bool _allocations_flag_prev_val##n = _memory_tracker_always_throw_logical_error_on_allocation; \ + _memory_tracker_always_throw_logical_error_on_allocation = val; \ + SCOPE_EXIT({ _memory_tracker_always_throw_logical_error_on_allocation = _allocations_flag_prev_val##n; }) +#define ALLOCATIONS_IN_SCOPE_IMPL(n, val) ALLOCATIONS_IN_SCOPE_IMPL_CONCAT(n, val) +#define DENY_ALLOCATIONS_IN_SCOPE ALLOCATIONS_IN_SCOPE_IMPL(__LINE__, true) +#define ALLOW_ALLOCATIONS_IN_SCOPE ALLOCATIONS_IN_SCOPE_IMPL(__LINE__, false) +#else +#define DENY_ALLOCATIONS_IN_SCOPE static_assert(true) +#define ALLOW_ALLOCATIONS_IN_SCOPE static_assert(true) +#endif /** Tracks memory consumption. * It throws an exception if amount of consumed memory become greater than certain limit. diff --git a/src/Common/QueryProfiler.cpp b/src/Common/QueryProfiler.cpp index 504d884dce0..bd1cab42be3 100644 --- a/src/Common/QueryProfiler.cpp +++ b/src/Common/QueryProfiler.cpp @@ -181,6 +181,7 @@ QueryProfilerReal::QueryProfilerReal(const UInt64 thread_id, const UInt32 period void QueryProfilerReal::signalHandler(int sig, siginfo_t * info, void * context) { + DENY_ALLOCATIONS_IN_SCOPE; writeTraceInfo(TraceType::Real, sig, info, context); } @@ -190,6 +191,7 @@ QueryProfilerCpu::QueryProfilerCpu(const UInt64 thread_id, const UInt32 period) void QueryProfilerCpu::signalHandler(int sig, siginfo_t * info, void * context) { + DENY_ALLOCATIONS_IN_SCOPE; writeTraceInfo(TraceType::CPU, sig, info, context); } diff --git a/src/Common/ThreadFuzzer.cpp b/src/Common/ThreadFuzzer.cpp index 88ff53534e6..a538ba7a49a 100644 --- a/src/Common/ThreadFuzzer.cpp +++ b/src/Common/ThreadFuzzer.cpp @@ -197,6 +197,7 @@ static void injection( void ThreadFuzzer::signalHandler(int) { + DENY_ALLOCATIONS_IN_SCOPE; auto saved_errno = errno; auto & fuzzer = ThreadFuzzer::instance(); diff --git a/src/Common/ThreadPool.cpp b/src/Common/ThreadPool.cpp index 7fc0d65aa5b..4b6834bd235 100644 --- a/src/Common/ThreadPool.cpp +++ b/src/Common/ThreadPool.cpp @@ -208,6 +208,7 @@ size_t ThreadPoolImpl::active() const template void ThreadPoolImpl::worker(typename std::list::iterator thread_it) { + DENY_ALLOCATIONS_IN_SCOPE; CurrentMetrics::Increment metric_all_threads( std::is_same_v ? CurrentMetrics::GlobalThread : CurrentMetrics::LocalThread); @@ -223,7 +224,7 @@ void ThreadPoolImpl::worker(typename std::list::iterator thread_ if (!jobs.empty()) { - job = std::move(jobs.top().job); + job = std::move(const_cast(jobs.top().job)); jobs.pop(); } else @@ -237,6 +238,7 @@ void ThreadPoolImpl::worker(typename std::list::iterator thread_ { try { + ALLOW_ALLOCATIONS_IN_SCOPE; CurrentMetrics::Increment metric_active_threads( std::is_same_v ? CurrentMetrics::GlobalThreadActive : CurrentMetrics::LocalThreadActive); diff --git a/src/Storages/System/StorageSystemStackTrace.cpp b/src/Storages/System/StorageSystemStackTrace.cpp index 0b5e82a1f3d..abb2fdf54ed 100644 --- a/src/Storages/System/StorageSystemStackTrace.cpp +++ b/src/Storages/System/StorageSystemStackTrace.cpp @@ -60,6 +60,7 @@ namespace void signalHandler(int, siginfo_t * info, void * context) { + DENY_ALLOCATIONS_IN_SCOPE; auto saved_errno = errno; /// We must restore previous value of errno in signal handler. /// In case malicious user is sending signals manually (for unknown reason). From 73e96250a802b18e9c2339945369bb5239c3d154 Mon Sep 17 00:00:00 2001 From: Alexander Tokmakov Date: Wed, 13 Jan 2021 15:19:41 +0300 Subject: [PATCH 2/2] add comments --- src/Common/MemoryTracker.h | 6 +++++- src/Common/ThreadPool.cpp | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/Common/MemoryTracker.h b/src/Common/MemoryTracker.h index b67f9e368e2..9a2c3a399ea 100644 --- a/src/Common/MemoryTracker.h +++ b/src/Common/MemoryTracker.h @@ -5,10 +5,14 @@ #include #include -#if !defined(NDEBUG) || defined(ADDRESS_SANITIZER) || defined(THREAD_SANITIZER) || defined(MEMORY_SANITIZER) || defined(UNDEFINED_BEHAVIOR_SANITIZER) +#if !defined(NDEBUG) #define MEMORY_TRACKER_DEBUG_CHECKS #endif +/// DENY_ALLOCATIONS_IN_SCOPE macro makes MemoryTracker throw LOGICAL_ERROR on any allocation attempt +/// until the end of the scope. It's useful to ensure that no allocations happen in signal handlers and +/// outside of try/catch block of thread functions. ALLOW_ALLOCATIONS_IN_SCOPE cancels effect of +/// DENY_ALLOCATIONS_IN_SCOPE in the inner scope. In Release builds these macros do nothing. #ifdef MEMORY_TRACKER_DEBUG_CHECKS #include extern thread_local bool _memory_tracker_always_throw_logical_error_on_allocation; diff --git a/src/Common/ThreadPool.cpp b/src/Common/ThreadPool.cpp index 4b6834bd235..92d02ae0f3e 100644 --- a/src/Common/ThreadPool.cpp +++ b/src/Common/ThreadPool.cpp @@ -224,6 +224,8 @@ void ThreadPoolImpl::worker(typename std::list::iterator thread_ if (!jobs.empty()) { + /// std::priority_queue does not provide interface for getting non-const reference to an element + /// to prevent us from modifying its priority. We have to use const_cast to force move semantics on JobWithPriority::job. job = std::move(const_cast(jobs.top().job)); jobs.pop(); }