Seal the heap: bounded library heaps and zero-allocation program tests

The programs point libuv, c-ares and Mbed TLS (MBEDTLS_PLATFORM_MEMORY in
the bundled build) at their own FixedHeap before anything runs, mark the
end of startup with heap_seal(), and in verbose mode report heap use and
any allocation after it at exit. The library build keeps the host's
allocator.

measure_alloc.py gains --udp, and the alloc_steady_programs_* tests now
fail on any malloc per relayed TCP connection (by name and by address)
or UDP association, for every cipher family.

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

src/core/heap_guard.cpp

@@ -0,0 +1,155 @@
+/* SPDX-License-Identifier: GPL-3.0-or-later */
+#include "core/heap_guard.hpp"
+
+#include <atomic>
+#include <cstdlib>
+#include <memory>
+
+#include <ares.h>
+#include <uv.h>
+
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+
+#include <mbedtls/platform.h>
+
+#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<std::size_t> 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 *>(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
+

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 <cstddef>
+#include <cstdint>
+
+#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

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;
 }
 

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;
 }

src/server.cpp

@@ -82,6 +82,7 @@ enum datatypes {
 #include <libnetfilter_conntrack/libnetfilter_conntrack_tcp.h>
 #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;
 }

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;
 }

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);

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);

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 $<TARGET_FILE:ss_alloc_audit> -- $<TARGET_FILE:test_cache_cpp>)
             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 $<TARGET_FILE:ss_alloc_audit>)
-            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 $<TARGET_FILE:ss_alloc_audit>
+                            --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 $<TARGET_FILE:ss_alloc_audit>
+                            --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

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)
 

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