diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b9393a22..16d1705b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -50,6 +50,7 @@ endif() # The C++ relay engine (src/relay) and the codecs it runs on. set(SS_RELAY_SOURCE + core/heap_guard.cpp relay/codec.cpp relay/socket.cpp relay/tcp.cpp diff --git a/src/core/heap_guard.cpp b/src/core/heap_guard.cpp new file mode 100644 index 00000000..ed7271bc --- /dev/null +++ b/src/core/heap_guard.cpp @@ -0,0 +1,155 @@ +/* SPDX-License-Identifier: GPL-3.0-or-later */ +#include "core/heap_guard.hpp" + +#include +#include +#include + +#include +#include + +#ifdef HAVE_CONFIG_H +#include "config.h" +#endif + +#include + +#include "core/fixed_heap.hpp" +#include "utils.h" + +namespace ss { + +namespace { + +struct Heaps { + FixedHeap uv, cares, mbedtls; +}; + +// Constructed in install_library_heaps() and never destroyed: libraries may +// free into their heap until the very end of the process. +alignas(Heaps) std::byte heaps_storage[sizeof(Heaps)]; +Heaps *heaps = nullptr; +std::atomic fallbacks{ 0 }; + +void *heap_malloc(FixedHeap &h, std::size_t size) noexcept +{ + if (void *p = h.allocate(size)) { + return p; + } + fallbacks.fetch_add(1, std::memory_order_relaxed); + return std::malloc(size); // heap-ok: over budget, see heap_report() +} + +void *heap_calloc(FixedHeap &h, std::size_t count, std::size_t size) noexcept +{ + if (void *p = h.allocate_zeroed(count, size)) { + return p; + } + fallbacks.fetch_add(1, std::memory_order_relaxed); + return std::calloc(count, size); // heap-ok: over budget, see heap_report() +} + +void *heap_realloc(FixedHeap &h, void *old, std::size_t size) noexcept +{ + if (old != nullptr && !h.owns(old)) { + return std::realloc(old, size); // heap-ok: allocated before the switch + } + if (void *p = h.reallocate(old, size)) { + return p; + } + if (size == 0) { + return nullptr; + } + fallbacks.fetch_add(1, std::memory_order_relaxed); + return nullptr; // the caller keeps `old` +} + +void heap_free(FixedHeap &h, void *p) noexcept +{ + if (p == nullptr) { + return; + } + if (h.owns(p)) { + h.deallocate(p); + } else { + std::free(p); // heap-ok: allocated before the switch or as a fallback + } +} + +void *uv_malloc_fn(std::size_t size) { return heap_malloc(heaps->uv, size); } +void *uv_realloc_fn(void *p, std::size_t size) { return heap_realloc(heaps->uv, p, size); } +void *uv_calloc_fn(std::size_t count, std::size_t size) { return heap_calloc(heaps->uv, count, size); } +void uv_free_fn(void *p) { heap_free(heaps->uv, p); } + +void *ares_malloc_fn(std::size_t size) { return heap_malloc(heaps->cares, size); } +void *ares_realloc_fn(void *p, std::size_t size) { return heap_realloc(heaps->cares, p, size); } +void ares_free_fn(void *p) { heap_free(heaps->cares, p); } + +#ifdef MBEDTLS_PLATFORM_MEMORY +void *mbedtls_calloc_fn(std::size_t count, std::size_t size) { return heap_calloc(heaps->mbedtls, count, size); } +void mbedtls_free_fn(void *p) { heap_free(heaps->mbedtls, p); } +#endif + +} // namespace + +Status install_library_heaps(const LibraryHeapBudgets &budgets) noexcept +{ + if (heaps != nullptr) { + return fail(Errc::invalid); + } + heaps = std::construct_at(reinterpret_cast(heaps_storage)); + SS_TRY(heaps->uv.reserve(budgets.uv)); + SS_TRY(heaps->cares.reserve(budgets.cares)); + SS_TRY(heaps->mbedtls.reserve(budgets.mbedtls)); + if (uv_replace_allocator(&uv_malloc_fn, &uv_realloc_fn, &uv_calloc_fn, &uv_free_fn) != 0) { + return fail(Errc::invalid); + } + if (ares_library_init_mem(ARES_LIB_INIT_ALL, &ares_malloc_fn, &ares_free_fn, &ares_realloc_fn) != ARES_SUCCESS) { + return fail(Errc::invalid); + } +#ifdef MBEDTLS_PLATFORM_MEMORY + if (mbedtls_platform_set_calloc_free(&mbedtls_calloc_fn, &mbedtls_free_fn) != 0) { + return fail(Errc::invalid); + } +#endif + return {}; +} + +void heap_seal() noexcept +{ + ss_heap_sealed = 1; +} + +bool heap_sealed() noexcept +{ + return ss_heap_sealed != 0; +} + +HeapReport heap_report() noexcept +{ + HeapReport r; + if (heaps != nullptr) { + auto uv = heaps->uv.stats(), ca = heaps->cares.stats(), mb = heaps->mbedtls.stats(); + r.uv_in_use = uv.in_use; + r.uv_peak = uv.peak; + r.cares_in_use = ca.in_use; + r.cares_peak = ca.peak; + r.mbedtls_in_use = mb.in_use; + r.mbedtls_peak = mb.peak; + } + r.fallbacks = fallbacks.load(std::memory_order_relaxed); + r.after_seal = ss_heap_late_allocations; + return r; +} + +void log_heap_report() noexcept +{ + HeapReport r = heap_report(); + LOGI("heaps: libuv %zu KiB (peak %zu), c-ares %zu KiB (peak %zu), mbedtls %zu KiB (peak %zu), " + "%zu fallback(s), %zu allocation(s) after startup", + r.uv_in_use / 1024, r.uv_peak / 1024, r.cares_in_use / 1024, r.cares_peak / 1024, r.mbedtls_in_use / 1024, + r.mbedtls_peak / 1024, r.fallbacks, r.after_seal); +} + +} // namespace ss + diff --git a/src/core/heap_guard.hpp b/src/core/heap_guard.hpp new file mode 100644 index 00000000..e54d5e5e --- /dev/null +++ b/src/core/heap_guard.hpp @@ -0,0 +1,50 @@ +/* SPDX-License-Identifier: GPL-3.0-or-later */ +#ifndef SS_CORE_HEAP_GUARD_HPP +#define SS_CORE_HEAP_GUARD_HPP + +// Bounded heaps for the third-party libraries, and the "no heap after +// startup" seal. +// +// install_library_heaps() points libuv, c-ares and (when built with +// MBEDTLS_PLATFORM_MEMORY) Mbed TLS at their own FixedHeap, reserved once. +// Programs call it first thing in main(); the library build never does, so +// an application embedding libshadowsocks-c keeps its allocator. Memory a +// library got before the switch is still freed through the system +// allocator, and a request no heap can serve falls back to it as well, +// counted as a fallback. +// +// heap_seal() marks the end of startup. After it, ss_malloc() and friends +// count every call (heap_report().after_seal); verbose mode logs the +// report at exit. + +#include +#include + +#include "core/result.hpp" + +namespace ss { + +struct LibraryHeapBudgets { + std::size_t uv = std::size_t{ 4 } << 20; + std::size_t cares = std::size_t{ 8 } << 20; + std::size_t mbedtls = std::size_t{ 1 } << 20; +}; + +struct HeapReport { + std::size_t uv_in_use = 0, uv_peak = 0; + std::size_t cares_in_use = 0, cares_peak = 0; + std::size_t mbedtls_in_use = 0, mbedtls_peak = 0; + std::size_t fallbacks = 0; // requests served by the system allocator + std::size_t after_seal = 0; // ss_malloc() calls after heap_seal() +}; + +Status install_library_heaps(const LibraryHeapBudgets &budgets = {}) noexcept; +void heap_seal() noexcept; +bool heap_sealed() noexcept; +HeapReport heap_report() noexcept; +// Log the report (verbose mode does this at exit). +void log_heap_report() noexcept; + +} // namespace ss + +#endif diff --git a/src/local.cpp b/src/local.cpp index 2147f23d..dd71cd34 100644 --- a/src/local.cpp +++ b/src/local.cpp @@ -64,6 +64,7 @@ #include "common.h" #include "ss_windows.h" +#include "core/heap_guard.hpp" #include "relay/tcp.hpp" #include "relay/udp.hpp" @@ -447,6 +448,10 @@ plugin_watcher_cb(SS_P_ ss_io *w, int revents) int main(int argc, char **argv) { + // Bounded heaps for libuv, c-ares and Mbed TLS before any of them runs. + if (!ss::install_library_heaps()) + FATAL("failed to reserve the library heaps"); + int i, c; int pid_flags = 0; int mtu = 0; @@ -1274,6 +1279,7 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`. #endif // Enter the loop + ss::heap_seal(); ss_run(loop, 0); if (verbose) { @@ -1296,6 +1302,9 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`. winsock_cleanup(); #endif + if (verbose) + ss::log_heap_report(); + return ret_val; } diff --git a/src/redir.cpp b/src/redir.cpp index 5e46529d..6c8d9d2a 100644 --- a/src/redir.cpp +++ b/src/redir.cpp @@ -56,6 +56,7 @@ #include "jconf.h" #include "common.h" +#include "core/heap_guard.hpp" #include "relay/tcp.hpp" #include "relay/udp.hpp" @@ -224,6 +225,10 @@ signal_cb(SS_P_ ss_signal *w, int revents) int main(int argc, char **argv) { + // Bounded heaps for libuv, c-ares and Mbed TLS before any of them runs. + if (!ss::install_library_heaps()) + FATAL("failed to reserve the library heaps"); + int i, c; int pid_flags = 0; int mptcp = 0; @@ -944,6 +949,7 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`. LOGI("running from root user"); } + ss::heap_seal(); ss_run(loop, 0); if (plugin != NULL) { @@ -957,5 +963,8 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`. } (void)tcp_running; + if (verbose) + ss::log_heap_report(); + return ret_val; } diff --git a/src/server.cpp b/src/server.cpp index 1f637cba..d5acf1ae 100644 --- a/src/server.cpp +++ b/src/server.cpp @@ -82,6 +82,7 @@ enum datatypes { #include #endif +#include "core/heap_guard.hpp" #include "relay/tcp.hpp" #include "relay/udp.hpp" @@ -821,6 +822,10 @@ plugin_watcher_cb(SS_P_ ss_io *w, int revents) int main(int argc, char **argv) { + // Bounded heaps for libuv, c-ares and Mbed TLS before any of them runs. + if (!ss::install_library_heaps()) + FATAL("failed to reserve the library heaps"); + int i, c; int pid_flags = 0; int mptcp = 0; @@ -1678,6 +1683,7 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`. #endif // start ev loop + ss::heap_seal(); ss_run(loop, 0); if (verbose) { @@ -1719,5 +1725,8 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`. winsock_cleanup(); #endif + if (verbose) + ss::log_heap_report(); + return ret_val; } diff --git a/src/tunnel.cpp b/src/tunnel.cpp index 67a924de..30c3b7a8 100644 --- a/src/tunnel.cpp +++ b/src/tunnel.cpp @@ -58,6 +58,7 @@ #include "common.h" #include "ss_windows.h" +#include "core/heap_guard.hpp" #include "relay/tcp.hpp" #include "relay/udp.hpp" @@ -246,6 +247,10 @@ plugin_watcher_cb(SS_P_ ss_io *w, int revents) int main(int argc, char **argv) { + // Bounded heaps for libuv, c-ares and Mbed TLS before any of them runs. + if (!ss::install_library_heaps()) + FATAL("failed to reserve the library heaps"); + int i, c; int pid_flags = 0; int mptcp = 0; @@ -1051,6 +1056,7 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`. } #endif + ss::heap_seal(); ss_run(loop, 0); if (plugin != NULL) { @@ -1071,5 +1077,8 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`. winsock_cleanup(); #endif + if (verbose) + ss::log_heap_report(); + return ret_val; } diff --git a/src/utils.cpp b/src/utils.cpp index 24052734..0b6d61ae 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -271,9 +271,21 @@ FATAL(const char *msg) exit(-1); } +/* Set once startup is over (see core/heap_guard.hpp); allocations through + * these helpers are counted from then on. */ +int ss_heap_sealed = 0; +size_t ss_heap_late_allocations = 0; + +#define NOTE_ALLOCATION() \ + do { \ + if (ss_heap_sealed) \ + ss_heap_late_allocations++; \ + } while (0) + void * ss_malloc(size_t size) { + NOTE_ALLOCATION(); void *tmp = malloc(size); if (tmp == NULL) exit(EXIT_FAILURE); @@ -283,6 +295,7 @@ ss_malloc(size_t size) void * ss_aligned_malloc(size_t size) { + NOTE_ALLOCATION(); int err; void *tmp = NULL; #ifdef HAVE_POSIX_MEMALIGN @@ -304,6 +317,7 @@ ss_aligned_malloc(size_t size) void * ss_realloc(void *ptr, size_t new_size) { + NOTE_ALLOCATION(); void *resized = realloc(ptr, new_size); if (resized == NULL) { free(ptr); diff --git a/src/utils.h b/src/utils.h index 1ffe1801..00aaa2b4 100644 --- a/src/utils.h +++ b/src/utils.h @@ -231,6 +231,8 @@ char *ss_strndup(const char *s, size_t n); int set_nofile(int nofile); #endif +extern int ss_heap_sealed; +extern size_t ss_heap_late_allocations; void *ss_malloc(size_t size); void *ss_aligned_malloc(size_t size); void *ss_realloc(void *ptr, size_t new_size); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c61f1ade..de229144 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -343,11 +343,33 @@ if(SS_BUILD_EXECUTABLES AND NOT CMAKE_CROSSCOMPILING) COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/alloc_audit/check_steady.py --shim $ -- $) set_tests_properties(alloc_steady_cache PROPERTIES LABELS "unit" TIMEOUT 60) - add_test(NAME alloc_audit_report - COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/alloc_audit/measure_alloc.py - --bin ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} - --shim $) - set_tests_properties(alloc_audit_report PROPERTIES LABELS "integration" TIMEOUT 120) + # The programs: no heap allocation per relayed connection once + # started, by name (DNS through the bounded c-ares heap) and by + # address, for each cipher family. + set(SS_ALLOC_CASES "aes-256-gcm@localhost" "2022-blake3-aes-256-gcm@127.0.0.1" + "chacha20-ietf-poly1305@127.0.0.1") + if(SS_ENABLE_LEGACY) + list(APPEND SS_ALLOC_CASES "aes-256-cfb@127.0.0.1") + endif() + foreach(case ${SS_ALLOC_CASES}) + string(REPLACE "@" ";" parts "${case}") + list(GET parts 0 method) + list(GET parts 1 target) + add_test(NAME alloc_steady_programs_${method} + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/alloc_audit/measure_alloc.py + --bin ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} + --shim $ + --method ${method} --target ${target} --max-per-connection 0) + set_tests_properties(alloc_steady_programs_${method} PROPERTIES LABELS "integration" TIMEOUT 120) + endforeach() + foreach(method aes-256-gcm 2022-blake3-aes-256-gcm) + add_test(NAME alloc_steady_programs_udp_${method} + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/alloc_audit/measure_alloc.py + --bin ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} + --shim $ + --method ${method} --udp --warmup 20 --connections 50 --max-per-connection 0) + set_tests_properties(alloc_steady_programs_udp_${method} PROPERTIES LABELS "integration" TIMEOUT 120) + endforeach() endif() add_test(NAME test_printable_salt_interop COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/interop.py diff --git a/tests/alloc_audit/measure_alloc.py b/tests/alloc_audit/measure_alloc.py index c23be554..4c6c8db3 100644 --- a/tests/alloc_audit/measure_alloc.py +++ b/tests/alloc_audit/measure_alloc.py @@ -27,7 +27,8 @@ import threading import time sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) -from interop import TCPHandler, TCPOrigin, free_port, peers, receive, socks # noqa: E402 +from interop import (TCPHandler, TCPOrigin, UDPHandler, UDPOrigin, free_port, peers, receive, socks, # noqa: E402 + udp_case) def read_snapshot(path, previous_lines, timeout=10): @@ -59,7 +60,12 @@ def tcp_case(local_port, target, origin_port, size): upload.result(timeout=10) -def relay(local_port, target, origin_port, connections, concurrency, size): +def relay(local_port, target, origin_port, connections, concurrency, size, udp=False): + if udp: + # Each case is a UDP association with four datagrams. + for _ in range(connections): + udp_case(local_port, origin_port) + return with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: futures = [pool.submit(tcp_case, local_port, target, origin_port, size) for _ in range(connections)] for future in futures: @@ -76,6 +82,7 @@ def main(): parser.add_argument("--concurrency", type=int, default=10) parser.add_argument("--size", type=int, default=16384, help="bytes echoed per connection") parser.add_argument("--max-per-connection", type=float, help="fail if any peer exceeds this") + parser.add_argument("--udp", action="store_true", help="relay UDP associations instead of TCP streams") parser.add_argument("--target", default="localhost", help="destination host sent to the server; a name makes the server resolve it") parser.add_argument("--json", action="store_true", help="print the result as JSON") @@ -87,17 +94,19 @@ def main(): key_size = 16 if args.method == "2022-blake3-aes-128-gcm" else 32 password = base64.b64encode(os.urandom(key_size)).decode("ascii") - with TCPOrigin(("127.0.0.1", 0), TCPHandler) as origin, \ + origin_class, handler = (UDPOrigin, UDPHandler) if args.udp else (TCPOrigin, TCPHandler) + with origin_class(("127.0.0.1", 0), handler) as origin, \ tempfile.TemporaryDirectory(prefix="ss-alloc-") as temp: threading.Thread(target=origin.serve_forever, daemon=True).start() server_port, local_port = free_port(), free_port() names = ["ss-server", "ss-local"] outputs = [Path(temp) / f"{name}.counts" for name in names] + udp = ["-u"] if args.udp else [] commands = [ [str(binary_dir / "ss-server"), "-s", "127.0.0.1", "-p", str(server_port), - "-k", password, "-m", args.method], + "-k", password, "-m", args.method] + udp, [str(binary_dir / "ss-local"), "-s", "127.0.0.1", "-p", str(server_port), - "-l", str(local_port), "-k", password, "-m", args.method], + "-l", str(local_port), "-k", password, "-m", args.method] + udp, ] environments = [] for output in outputs: @@ -112,10 +121,11 @@ def main(): peers([commands[1]], [local_port], environments[1]) as local: processes = server + local counts = [0, 0] - relay(local_port, args.target, origin.server_address[1], args.warmup, args.concurrency, args.size) + relay(local_port, args.target, origin.server_address[1], args.warmup, args.concurrency, args.size, args.udp) time.sleep(.5) before = snapshot_all(processes, outputs, counts) - relay(local_port, args.target, origin.server_address[1], args.connections, args.concurrency, args.size) + relay(local_port, args.target, origin.server_address[1], args.connections, args.concurrency, args.size, + args.udp) time.sleep(.5) after = snapshot_all(processes, outputs, counts) diff --git a/third_party/mbedtls_config.h b/third_party/mbedtls_config.h index c03c7cb5..bfe68e65 100644 --- a/third_party/mbedtls_config.h +++ b/third_party/mbedtls_config.h @@ -19,6 +19,8 @@ #define MBEDTLS_SHA256_C #define MBEDTLS_SHA512_C #define MBEDTLS_PLATFORM_C +/* The programs route Mbed TLS allocations into a bounded heap. */ +#define MBEDTLS_PLATFORM_MEMORY #define MBEDTLS_VERSION_C #define MBEDTLS_VERSION_FEATURES #define MBEDTLS_CIPHER_MODE_CTR