C++20 port 12/14: seal the heap, capacity limits (#3078)

* 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.

* Fix GCC -Wsign-compare errors in the converted C modules

GCC enables -Wsign-compare with -Wall for C++ (clang does not), which
turned a few signed/unsigned comparisons in base64, netutils, server and
manager into errors; udp.cpp needed <ctime>.

* Add --max-connections and --max-udp-sessions

The TCP connection pool and the UDP session table are sized from these
options (or max_connections / max_udp_sessions in the configuration
file) at startup instead of fixed constants; the DNS query pool follows
the connection limit. test_limits.py checks that a connection beyond
the limit is refused and that accepting resumes once one closes.

doc/shadowsocks-c.md

@@ -79,6 +79,8 @@ The config file equivalent of command line options is listed as examples below.
 - `--fast-open`: `"fast_open": true`
 - `--reuse-port`: `"reuse_port": true`
 - `--no-delay`: `"no_delay": true`
+- `--max-connections 1024`: `"max_connections": 1024`
+- `--max-udp-sessions 512`: `"max_udp_sessions": 512`
 - `--printable-salt`: `"printable_salt": true` (ss-local only)
 - `--plugin "obfs-server"`: `"plugin": "obfs-server"`
 - `--plugin-opts "obfs=http"`: `"plugin_opts": "obfs=http"`

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/base64.cpp

@@ -91,7 +91,7 @@ char *base64_encode(char *out, int out_size, const uint8_t *in, int in_size)
     int i_shift = 0;
     int bytes_remaining = in_size;
 
-    if (in_size >= UINT_MAX / 4 ||
+    if ((unsigned)in_size >= UINT_MAX / 4 ||
         out_size < BASE64_SIZE(in_size))
         return NULL;
     ret = dst = out;

src/common.h

@@ -65,6 +65,8 @@ enum {
     GETOPT_VAL_REUSE_PORT,
     GETOPT_VAL_FAST_OPEN,
     GETOPT_VAL_NODELAY,
+    GETOPT_VAL_MAX_CONNECTIONS,
+    GETOPT_VAL_MAX_UDP_SESSIONS,
     GETOPT_VAL_PRINTABLE_SALT,
     GETOPT_VAL_ACL,
     GETOPT_VAL_MTU,

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/jconf.cpp

@@ -389,6 +389,10 @@ read_jconf(const char *file)
                 conf.tcp_outgoing_rcvbuf = to_nonnegative_int(value, name);
             } else if (strcmp(name, "auth") == 0) {
                 FATAL("One time auth has been deprecated. Try AEAD ciphers instead.");
+            } else if (strcmp(name, "max_connections") == 0) {
+                conf.max_connections = to_nonnegative_int(value, name);
+            } else if (strcmp(name, "max_udp_sessions") == 0) {
+                conf.max_udp_sessions = to_nonnegative_int(value, name);
             } else if (strcmp(name, "nofile") == 0) {
                 conf.nofile = to_nonnegative_int(value, name);
             } else if (strcmp(name, "nameserver") == 0) {

src/jconf.h

@@ -92,6 +92,8 @@ typedef struct {
     int mptcp;
     int ipv6_first;
     int no_delay;
+    int max_connections;
+    int max_udp_sessions;
     int printable_salt;
     int tcp_tproxy;
     char *workdir;

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"
 
@@ -110,6 +111,8 @@ static int mode      = TCP_ONLY;
 static int ipv6first = 0;
 int fast_open        = 0;
 static int no_delay  = 0;
+static int max_connections  = 0;
+static int max_udp_sessions = 0;
 static int printable_salt = 0;
 static int udp_fd    = 0;
 static int ret_val   = 0;
@@ -322,7 +325,7 @@ struct Relays {
         config.kind         = ss::relay::Kind::local;
         config.remote       = *remote;
         config.timeout_secs = timeout;
-        config.max_sessions = 256;
+        config.max_sessions = max_udp_sessions > 0 ? max_udp_sessions : 256;
         config.mtu          = mtu;
         config.iface        = iface;
 #ifdef __ANDROID__
@@ -371,7 +374,7 @@ fill_tcp_config(ss::relay::TcpConfig &config, const struct sockaddr_storage *rem
     config.kind             = ss::relay::Kind::local;
     config.remotes          = { remotes, (size_t)remote_count };
     config.timeout_secs     = timeout;
-    config.max_connections  = DEFAULT_MAX_CONNECTIONS;
+    config.max_connections  = max_connections > 0 ? max_connections : DEFAULT_MAX_CONNECTIONS;
     config.fast_open        = fast_open != 0;
     config.no_delay         = no_delay != 0;
     config.printable_salt   = printable_salt != 0;
@@ -447,6 +450,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;
@@ -611,6 +618,8 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`.
 \snippet{doc} utils.cpp cli_long_tcp_outgoing_rcvbuf
 \snippet{doc} utils.cpp cli_long_fast_open
 \snippet{doc} utils.cpp cli_long_no_delay
+\snippet{doc} utils.cpp cli_long_max_connections
+\snippet{doc} utils.cpp cli_long_max_udp_sessions
 \snippet{doc} utils.cpp cli_long_acl
 \snippet{doc} utils.cpp cli_long_mtu
 \snippet{doc} utils.cpp cli_long_mptcp
@@ -653,6 +662,8 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`.
         { "fast-open",   no_argument,       NULL, GETOPT_VAL_FAST_OPEN   },
         { "printable-salt", no_argument, NULL, GETOPT_VAL_PRINTABLE_SALT },
         { "no-delay",    no_argument,       NULL, GETOPT_VAL_NODELAY     },
+        { "max-connections", required_argument, NULL, GETOPT_VAL_MAX_CONNECTIONS },
+        { "max-udp-sessions", required_argument, NULL, GETOPT_VAL_MAX_UDP_SESSIONS },
         { "acl",         required_argument, NULL, GETOPT_VAL_ACL         },
         { "mtu",         required_argument, NULL, GETOPT_VAL_MTU         },
         { "mptcp",       no_argument,       NULL, GETOPT_VAL_MPTCP       },
@@ -711,6 +722,16 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`.
         case GETOPT_VAL_PRINTABLE_SALT:
             printable_salt = 1;
             break;
+        case GETOPT_VAL_MAX_CONNECTIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_connections) == -1) {
+                cli_error("invalid connection limit", c, NULL);
+            }
+            break;
+        case GETOPT_VAL_MAX_UDP_SESSIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_udp_sessions) == -1) {
+                cli_error("invalid UDP session limit", c, NULL);
+            }
+            break;
         case GETOPT_VAL_NODELAY:
             no_delay = 1;
             LOGI("enable TCP no-delay");
@@ -952,6 +973,12 @@ Same behavior as `-S`; see that option for details. Short alias: `-S`.
         if (no_delay == 0) {
             no_delay = conf->no_delay;
         }
+        if (max_connections == 0) {
+            max_connections = conf->max_connections;
+        }
+        if (max_udp_sessions == 0) {
+            max_udp_sessions = conf->max_udp_sessions;
+        }
         if (printable_salt == 0) {
             printable_salt = conf->printable_salt;
         }
@@ -1274,6 +1301,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 +1324,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/manager.cpp

@@ -605,7 +605,7 @@ get_server(char *buf, int len)
     int valid = 1;
     if (obj->type == json_object) {
         int i = 0;
-        for (i = 0; i < obj->u.object.length; i++) {
+        for (i = 0; (unsigned int)i < obj->u.object.length; i++) {
             char *name        = obj->u.object.values[i].name;
             json_value *value = obj->u.object.values[i].value;
             if (strcmp(name, "server_port") == 0) {
@@ -699,7 +699,7 @@ parse_traffic(char *buf, int len, char *port, uint64_t *traffic)
 
     if (obj->type == json_object) {
         int i = 0;
-        for (i = 0; i < obj->u.object.length; i++) {
+        for (i = 0; (unsigned int)i < obj->u.object.length; i++) {
             char *name        = obj->u.object.values[i].name;
             json_value *value = obj->u.object.values[i].value;
             if (value->type == json_integer) {
@@ -1032,7 +1032,7 @@ manager_recv_cb(SS_P_ ss_io *w, int revents)
             }
             if (pos >= BUF_SIZE - entry_pos) {
                 if (sendto(manager->fd, buf, pos, 0, (struct sockaddr *)&claddr, len)
-                    != pos) {
+                    != (ssize_t)pos) {
                     ERROR("list_sendto");
                 }
                 memset(buf, 0, BUF_SIZE);
@@ -1045,7 +1045,7 @@ manager_recv_cb(SS_P_ ss_io *w, int revents)
         strcpy(buf + (pos > 2 ? pos - 1 : 1), "\n]"); // Remove trailing ","
         pos = strlen(buf);
         if (sendto(manager->fd, buf, pos, 0, (struct sockaddr *)&claddr, len)
-            != pos) {
+            != (ssize_t)pos) {
             ERROR("list_sendto");
         }
     } else if (strcmp(action, "remove") == 0) {
@@ -1107,7 +1107,7 @@ manager_recv_cb(SS_P_ ss_io *w, int revents)
                     goto ERROR_MSG;
                 }
                 if (sendto(manager->fd, buf, pos, 0, (struct sockaddr *)&claddr, len)
-                    != pos) {
+                    != (ssize_t)pos) {
                     ERROR("ping_sendto");
                 }
                 memset(buf, 0, BUF_SIZE);
@@ -1132,7 +1132,7 @@ manager_recv_cb(SS_P_ ss_io *w, int revents)
         }
 
         if (sendto(manager->fd, buf, pos, 0, (struct sockaddr *)&claddr, len)
-            != pos) {
+            != (ssize_t)pos) {
             ERROR("ping_sendto");
         }
     }

src/netutils.cpp

@@ -341,7 +341,7 @@ validate_hostname(const char *hostname, const int hostname_len)
 int
 is_ipv6only(ss_addr_t *servers, size_t server_num, int ipv6first)
 {
-    int i;
+    size_t i;
     for (i = 0; i < server_num; i++) {
         struct sockaddr_storage storage;
         memset(&storage, 0, sizeof(struct sockaddr_storage));

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"
 
@@ -97,6 +98,8 @@ static int nofile = 0;
 #endif
 int fast_open       = 0;
 static int no_delay = 0;
+static int max_connections  = 0;
+static int max_udp_sessions = 0;
 static int ret_val  = 0;
 
 static struct ss_signal sigint_watcher;
@@ -224,6 +227,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;
@@ -382,6 +389,8 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
 \snippet{doc} utils.cpp cli_long_tcp_outgoing_sndbuf
 \snippet{doc} utils.cpp cli_long_tcp_outgoing_rcvbuf
 \snippet{doc} utils.cpp cli_long_no_delay
+\snippet{doc} utils.cpp cli_long_max_connections
+\snippet{doc} utils.cpp cli_long_max_udp_sessions
 \snippet{doc} utils.cpp cli_long_password
 \snippet{doc} utils.cpp cli_long_key
 \snippet{doc} utils.cpp cli_long_help
@@ -416,6 +425,8 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
         { "tcp-outgoing-sndbuf", required_argument, NULL, GETOPT_VAL_TCP_OUTGOING_SNDBUF },
         { "tcp-outgoing-rcvbuf", required_argument, NULL, GETOPT_VAL_TCP_OUTGOING_RCVBUF },
         { "no-delay",    no_argument,       NULL, GETOPT_VAL_NODELAY     },
+        { "max-connections", required_argument, NULL, GETOPT_VAL_MAX_CONNECTIONS },
+        { "max-udp-sessions", required_argument, NULL, GETOPT_VAL_MAX_UDP_SESSIONS },
         { "password",    required_argument, NULL, GETOPT_VAL_PASSWORD    },
         { "key",         required_argument, NULL, GETOPT_VAL_KEY         },
         { "help",        no_argument,       NULL, GETOPT_VAL_HELP        },
@@ -456,6 +467,16 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
             if (mptcp)
                 LOGI("enable multipath TCP (%s)", mptcp > 0 ? "out-of-tree" : "upstream");
             break;
+        case GETOPT_VAL_MAX_CONNECTIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_connections) == -1) {
+                cli_error("invalid connection limit", c, NULL);
+            }
+            break;
+        case GETOPT_VAL_MAX_UDP_SESSIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_udp_sessions) == -1) {
+                cli_error("invalid UDP session limit", c, NULL);
+            }
+            break;
         case GETOPT_VAL_NODELAY:
             no_delay = 1;
             LOGI("enable TCP no-delay");
@@ -651,6 +672,12 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
         if (no_delay == 0) {
             no_delay = conf->no_delay;
         }
+        if (max_connections == 0) {
+            max_connections = conf->max_connections;
+        }
+        if (max_udp_sessions == 0) {
+            max_udp_sessions = conf->max_udp_sessions;
+        }
         if (reuse_port == 0) {
             reuse_port = conf->reuse_port;
         }
@@ -861,7 +888,7 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
         ss::relay::UdpConfig udp_config;
         udp_config.kind         = ss::relay::Kind::redir;
         udp_config.timeout_secs = timeout_secs;
-        udp_config.max_sessions = 256;
+        udp_config.max_sessions = max_udp_sessions > 0 ? max_udp_sessions : 256;
         udp_config.mtu          = mtu;
         udp_config.reuse_port   = reuse_port != 0;
         char *host = remote_addr[0].host;
@@ -878,7 +905,7 @@ Same behavior as `-T`; see that option for details. Short alias: `-T`.
         tcp_config.kind                     = ss::relay::Kind::redir;
         tcp_config.remotes                  = { remote_storage, (size_t)remote_count };
         tcp_config.timeout_secs             = timeout_secs;
-        tcp_config.max_connections          = DEFAULT_MAX_CONNECTIONS;
+        tcp_config.max_connections          = max_connections > 0 ? max_connections : DEFAULT_MAX_CONNECTIONS;
         tcp_config.fast_open                = fast_open != 0;
         tcp_config.no_delay                 = no_delay != 0;
         tcp_config.tproxy                   = tcp_tproxy != 0;
@@ -944,6 +971,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 +985,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/relay/tcp.cpp

@@ -142,6 +142,7 @@ Status TcpRelay::init(event::Loop &loop, Method &method, const TcpConfig &config
     idle_list_.set_timeout(std::max<std::uint64_t>(MIN_TCP_IDLE_TIMEOUT, timeout) * 1000);
     SS_TRY(timer_.open(loop, *this));
     own_closed_ = false;
+    LOGI("tcp: up to %u connections (%zu KiB each)", config.max_connections, sizeof(Connection) / 1024);
     return {};
 }
 

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"
 
@@ -131,6 +132,8 @@ static int mode      = TCP_ONLY;
 static int ipv6first = 0;
 int fast_open        = 0;
 static int no_delay  = 0;
+static int max_connections  = 0;
+static int max_udp_sessions = 0;
 static int ret_val   = 0;
 
 #ifdef HAVE_SETRLIMIT
@@ -225,7 +228,7 @@ stat_update_cb(SS_P_ ss_timer *watcher, int revents)
         strncpy(svaddr.sun_path, manager_addr, sizeof(svaddr.sun_path) - 1);
 
         if (sendto(sfd, resp, strlen(resp) + 1, 0, (struct sockaddr *)&svaddr,
-                   sizeof(struct sockaddr_un)) != msgLen) {
+                   sizeof(struct sockaddr_un)) != (ssize_t)msgLen) {
             ERROR("stat_sendto");
             ss_socket_close(sfd);
             return;
@@ -821,6 +824,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;
@@ -982,6 +989,8 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
 \snippet{doc} utils.cpp cli_long_tcp_outgoing_sndbuf
 \snippet{doc} utils.cpp cli_long_tcp_outgoing_rcvbuf
 \snippet{doc} utils.cpp cli_long_no_delay
+\snippet{doc} utils.cpp cli_long_max_connections
+\snippet{doc} utils.cpp cli_long_max_udp_sessions
 \snippet{doc} utils.cpp cli_long_acl
 \snippet{doc} utils.cpp cli_long_manager_address
 \snippet{doc} utils.cpp cli_long_mtu
@@ -1019,6 +1028,8 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
         { "tcp-outgoing-sndbuf", required_argument, NULL, GETOPT_VAL_TCP_OUTGOING_SNDBUF },
         { "tcp-outgoing-rcvbuf", required_argument, NULL, GETOPT_VAL_TCP_OUTGOING_RCVBUF },
         { "no-delay",        no_argument,       NULL, GETOPT_VAL_NODELAY     },
+        { "max-connections", required_argument, NULL, GETOPT_VAL_MAX_CONNECTIONS },
+        { "max-udp-sessions", required_argument, NULL, GETOPT_VAL_MAX_UDP_SESSIONS },
         { "acl",             required_argument, NULL, GETOPT_VAL_ACL         },
         { "manager-address", required_argument, NULL,
           GETOPT_VAL_MANAGER_ADDRESS },
@@ -1060,6 +1071,16 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
         case GETOPT_VAL_FAST_OPEN:
             fast_open = 1;
             break;
+        case GETOPT_VAL_MAX_CONNECTIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_connections) == -1) {
+                cli_error("invalid connection limit", c, NULL);
+            }
+            break;
+        case GETOPT_VAL_MAX_UDP_SESSIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_udp_sessions) == -1) {
+                cli_error("invalid UDP session limit", c, NULL);
+            }
+            break;
         case GETOPT_VAL_NODELAY:
             no_delay = 1;
             LOGI("enable TCP no-delay");
@@ -1260,6 +1281,12 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
         if (no_delay == 0) {
             no_delay = conf->no_delay;
         }
+        if (max_connections == 0) {
+            max_connections = conf->max_connections;
+        }
+        if (max_udp_sessions == 0) {
+            max_udp_sessions = conf->max_udp_sessions;
+        }
         if (reuse_port == 0) {
             reuse_port = conf->reuse_port;
         }
@@ -1531,14 +1558,15 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
     ss::relay::UdpRelay udp_relay;
     bool tcp_running = false;
     if (!tcp_loop.adopt(ss_native_loop(loop))
-        || !tcp_resolver.init(tcp_loop, nameservers, ipv6first != 0, DEFAULT_MAX_CONNECTIONS + 64)) {
+        || !tcp_resolver.init(tcp_loop, nameservers, ipv6first != 0,
+                           (max_connections > 0 ? max_connections : DEFAULT_MAX_CONNECTIONS) + 64)) {
         FATAL("failed to set up the resolver");
     }
     if (mode != UDP_ONLY) {
         ss::relay::TcpConfig tcp_config;
         tcp_config.kind            = ss::relay::Kind::server;
         tcp_config.timeout_secs    = timeout_secs;
-        tcp_config.max_connections = DEFAULT_MAX_CONNECTIONS;
+        tcp_config.max_connections = max_connections > 0 ? max_connections : DEFAULT_MAX_CONNECTIONS;
         tcp_config.fast_open       = fast_open != 0;
         tcp_config.no_delay        = no_delay != 0;
         tcp_config.acl             = acl != 0;
@@ -1613,7 +1641,7 @@ Same behavior as `-d`; see that option for details. Short alias: `-d`.
         ss::relay::UdpConfig udp_config;
         udp_config.kind         = ss::relay::Kind::server;
         udp_config.timeout_secs = timeout_secs;
-        udp_config.max_sessions = 512;
+        udp_config.max_sessions = max_udp_sessions > 0 ? max_udp_sessions : 512;
         udp_config.mtu          = mtu;
         udp_config.iface        = iface;
         udp_config.label        = remote_port != NULL ? remote_port : "";
@@ -1678,6 +1706,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 +1748,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"
 
@@ -112,6 +113,8 @@ static int mode      = TCP_ONLY;
 static int nofile = 0;
 #endif
 static int no_delay = 0;
+static int max_connections  = 0;
+static int max_udp_sessions = 0;
 int fast_open       = 0;
 static int ret_val  = 0;
 
@@ -246,6 +249,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;
@@ -406,6 +413,8 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
 \snippet{doc} utils.cpp cli_long_fast_open
 \snippet{doc} utils.cpp cli_long_mtu
 \snippet{doc} utils.cpp cli_long_no_delay
+\snippet{doc} utils.cpp cli_long_max_connections
+\snippet{doc} utils.cpp cli_long_max_udp_sessions
 \snippet{doc} utils.cpp cli_long_mptcp
 \snippet{doc} utils.cpp cli_long_plugin
 \snippet{doc} utils.cpp cli_long_plugin_opts
@@ -444,6 +453,8 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
         { "fast-open",   no_argument,       NULL, GETOPT_VAL_FAST_OPEN   },
         { "mtu",         required_argument, NULL, GETOPT_VAL_MTU         },
         { "no-delay",    no_argument,       NULL, GETOPT_VAL_NODELAY     },
+        { "max-connections", required_argument, NULL, GETOPT_VAL_MAX_CONNECTIONS },
+        { "max-udp-sessions", required_argument, NULL, GETOPT_VAL_MAX_UDP_SESSIONS },
         { "mptcp",       no_argument,       NULL, GETOPT_VAL_MPTCP       },
         { "plugin",      required_argument, NULL, GETOPT_VAL_PLUGIN      },
         { "plugin-opts", required_argument, NULL, GETOPT_VAL_PLUGIN_OPTS },
@@ -497,6 +508,16 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
             if (mptcp)
                 LOGI("enable multipath TCP (%s)", mptcp > 0 ? "out-of-tree" : "upstream");
             break;
+        case GETOPT_VAL_MAX_CONNECTIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_connections) == -1) {
+                cli_error("invalid connection limit", c, NULL);
+            }
+            break;
+        case GETOPT_VAL_MAX_UDP_SESSIONS:
+            if (ss_parse_int(optarg, 1, 1 << 20, &max_udp_sessions) == -1) {
+                cli_error("invalid UDP session limit", c, NULL);
+            }
+            break;
         case GETOPT_VAL_NODELAY:
             no_delay = 1;
             LOGI("enable TCP no-delay");
@@ -703,6 +724,12 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
         if (no_delay == 0) {
             no_delay = conf->no_delay;
         }
+        if (max_connections == 0) {
+            max_connections = conf->max_connections;
+        }
+        if (max_udp_sessions == 0) {
+            max_udp_sessions = conf->max_udp_sessions;
+        }
         if (reuse_port == 0) {
             reuse_port = conf->reuse_port;
         }
@@ -960,7 +987,7 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
     tcp_config.kind              = ss::relay::Kind::tunnel;
     tcp_config.remotes           = { remote_storage, (size_t)remote_count };
     tcp_config.timeout_secs      = timeout_secs;
-    tcp_config.max_connections   = DEFAULT_MAX_CONNECTIONS;
+    tcp_config.max_connections   = max_connections > 0 ? max_connections : DEFAULT_MAX_CONNECTIONS;
     tcp_config.fast_open         = fast_open != 0;
     tcp_config.no_delay          = no_delay != 0;
     tcp_config.upstream.sndbuf   = tcp_outgoing_sndbuf;
@@ -1020,7 +1047,7 @@ Same behavior as `-V`; see that option for details. Short alias: `-V`.
         udp_config.tunnel_target     = tcp_config.tunnel_target;
         udp_config.tunnel_target_len = tcp_config.tunnel_target_len;
         udp_config.timeout_secs      = timeout_secs;
-        udp_config.max_sessions      = 256;
+        udp_config.max_sessions      = max_udp_sessions > 0 ? max_udp_sessions : 256;
         udp_config.mtu               = mtu;
         udp_config.iface             = iface;
 #ifdef __ANDROID__
@@ -1051,6 +1078,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 +1099,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);
@@ -471,6 +485,20 @@ Only available with MPTCP enabled Linux kernel.
 Enable TCP_NODELAY.
 [cli_long_no_delay] */
 
+/* [cli_long_max_connections]
+\par `--max-connections <number>`
+Serve at most this many TCP connections at once (default: 1024). Their
+memory is reserved at startup; a connection beyond the limit is refused
+until another one closes.
+[cli_long_max_connections] */
+
+/* [cli_long_max_udp_sessions]
+\par `--max-udp-sessions <number>`
+Keep at most this many UDP associations (default: 256 for clients, 512
+for ss-server). When the table is full, the least recently used one is
+dropped.
+[cli_long_max_udp_sessions] */
+
 /* [cli_long_tcp_incoming_sndbuf]
 \par `--tcp-incoming-sndbuf <size>`
 Set TCP send buffer size for incoming connections.
@@ -690,6 +718,10 @@ usage(void)
     cli_help_option("--fast-open", "Enable TCP Fast Open where supported.");
     cli_help_option("--reuse-port", "Enable port reuse where supported.");
     cli_help_option("--no-delay", "Enable TCP_NODELAY.");
+#ifndef MODULE_MANAGER
+    cli_help_option("--max-connections NUMBER", "TCP connections served at once (default: 1024).");
+    cli_help_option("--max-udp-sessions NUMBER", "UDP associations kept (default: 256; ss-server 512).");
+#endif
 #ifdef MODULE_LOCAL
     cli_help_option("--printable-salt", "Experimental printable TCP salt (chacha20-ietf-poly1305/aes-256-gcm only).");
 #endif

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
@@ -360,6 +382,10 @@ if(SS_BUILD_EXECUTABLES AND NOT CMAKE_CROSSCOMPILING)
         if(SS_ENABLE_LEGACY)
             list(APPEND SS_TUNNEL_INTEROP_ARGS --legacy)
         endif()
+        add_test(NAME test_limits
+            COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_limits.py
+                    --bin ${CMAKE_RUNTIME_OUTPUT_DIRECTORY})
+        set_tests_properties(test_limits PROPERTIES LABELS "integration" TIMEOUT 60)
         add_test(NAME test_tunnel_interop
             COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/tunnel_interop.py ${SS_TUNNEL_INTEROP_ARGS})
         set_tests_properties(test_tunnel_interop PROPERTIES LABELS "integration" TIMEOUT 300)

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)
 

tests/test_limits.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""--max-connections: connections beyond the limit are refused, and accepting
+resumes once a slot frees."""
+import argparse
+from pathlib import Path
+import socket
+import threading
+import time
+
+from interop import TCPHandler, TCPOrigin, free_port, peers, receive, socks
+
+
+def echo_once(sock):
+    sock.sendall(b"ping")
+    return receive(sock, 4) == b"ping"
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--bin", required=True, type=Path)
+    args = parser.parse_args()
+    bins = args.bin.resolve()
+    with TCPOrigin(("127.0.0.1", 0), TCPHandler) as origin:
+        threading.Thread(target=origin.serve_forever, daemon=True).start()
+        sport, lport = free_port(), free_port()
+        server = [str(bins / "ss-server"), "-s", "127.0.0.1", "-p", str(sport), "-k", "pw", "-m", "aes-256-gcm",
+                  "--max-connections", "2"]
+        local = [str(bins / "ss-local"), "-s", "127.0.0.1", "-p", str(sport), "-l", str(lport), "-k", "pw",
+                 "-m", "aes-256-gcm"]
+        with peers([server, local], [sport, lport]):
+            target = ("127.0.0.1", origin.server_address[1])
+            held = []
+            for _ in range(2):
+                sock, _ = socks(lport, 1, target)
+                assert echo_once(sock), "connection within the limit failed"
+                held.append(sock)
+            # The third is refused by ss-server, so ss-local drops it.
+            third, _ = socks(lport, 1, target)
+            third.settimeout(3)
+            try:
+                refused = not echo_once(third)
+            except (OSError, AssertionError):
+                refused = True
+            third.close()
+            assert refused, "a connection beyond --max-connections was served"
+            # Free a slot; accepting resumes.
+            held.pop().close()
+            deadline = time.monotonic() + 5
+            while True:
+                try:
+                    sock, _ = socks(lport, 1, target)
+                    sock.settimeout(3)
+                    if echo_once(sock):
+                        held.append(sock)
+                        break
+                except (OSError, AssertionError):
+                    pass
+                assert time.monotonic() < deadline, "accepting did not resume after a slot freed"
+                time.sleep(.1)
+            for sock in held:
+                sock.close()
+    print("PASS --max-connections refuses and recovers")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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