From 2e9a6353ff6b189ed4c046c4d2765c044b9cd593 Mon Sep 17 00:00:00 2001 From: Matt Corallo Date: Sun, 19 Jul 2020 17:17:03 -0400 Subject: [PATCH] Globablly limit alloc to 1.75GB --- src/main.rs | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/src/main.rs b/src/main.rs index c496f14..316ca12 100644 --- a/src/main.rs +++ b/src/main.rs @@ -40,6 +40,48 @@ static mut TOR_PROXY: Option = None; pub static START_SHUTDOWN: AtomicBool = AtomicBool::new(false); static SCANNING: AtomicBool = AtomicBool::new(false); + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::ptr; +use std::sync::atomic::AtomicUsize; + +// We keep track of all memory allocated by Rust code, refusing new allocations if it exceeds +// 1.75GB. +// +// Note that while Rust's std, in general, should panic in response to a null allocation, it +// is totally conceivable that some code will instead dereference this null pointer, which +// would violate our guarantees that Rust modules should never crash the entire application. +// +// In the future, as upstream Rust explores a safer allocation API (eg the Alloc API which +// returns Results instead of raw pointers, or redefining the GlobalAlloc API to allow +// panic!()s inside of alloc calls), we should switch to those, however these APIs are +// currently unstable. +const TOTAL_MEM_LIMIT_BYTES: usize = (1024 + 756) * 1024 * 1024; +static TOTAL_MEM_ALLOCD: AtomicUsize = AtomicUsize::new(0); +struct MemoryLimitingAllocator; +unsafe impl GlobalAlloc for MemoryLimitingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let len = layout.size(); + if len > TOTAL_MEM_LIMIT_BYTES { + return ptr::null_mut(); + } + if TOTAL_MEM_ALLOCD.fetch_add(len, Ordering::AcqRel) + len > TOTAL_MEM_LIMIT_BYTES { + TOTAL_MEM_ALLOCD.fetch_sub(len, Ordering::AcqRel); + return ptr::null_mut(); + } + System.alloc(layout) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + System.dealloc(ptr, layout); + TOTAL_MEM_ALLOCD.fetch_sub(layout.size(), Ordering::AcqRel); + } +} + +#[global_allocator] +static ALLOC: MemoryLimitingAllocator = MemoryLimitingAllocator; + + struct PeerState { request: Arc<(u64, BlockHash, Block)>, node_services: u64, -- 2.30.2