From 641643ec48a326f3dd43fa04658fb5593af874ca Mon Sep 17 00:00:00 2001 From: Max Lv Date: Wed, 14 Oct 2026 12:51:47 +0800 Subject: [PATCH] Add heap allocation audit for relayed connections alloc_counter is a preloadable malloc counter (glibc, macOS) that writes a snapshot on SIGUSR2. measure_alloc.py warms ss-local and ss-server up, snapshots both, relays more connections and reports allocations per connection. It runs report-only as the alloc_audit_report ctest. interop.peers() now yields the started processes. --- tests/alloc_audit/alloc_counter.c | 217 +++++++++++++++++++++++++++++ tests/alloc_audit/measure_alloc.py | 139 ++++++++++++++++++ tests/interop.py | 2 +- 3 files changed, 357 insertions(+), 1 deletion(-) create mode 100644 tests/alloc_audit/alloc_counter.c create mode 100644 tests/alloc_audit/measure_alloc.py diff --git a/tests/alloc_audit/alloc_counter.c b/tests/alloc_audit/alloc_counter.c new file mode 100644 index 00000000..c3ed5a70 --- /dev/null +++ b/tests/alloc_audit/alloc_counter.c @@ -0,0 +1,217 @@ +/* SPDX-License-Identifier: GPL-3.0-or-later */ +/* + * alloc_counter: a preloadable shim that counts heap allocations. + * + * Load it into a program with LD_PRELOAD (glibc) or DYLD_INSERT_LIBRARIES + * (macOS). Every SIGUSR2 appends one line with the running totals to the + * file named by SS_ALLOC_AUDIT_OUT: + * + * + * + * A driver takes a snapshot after warming a proxy up and another after more + * traffic; the difference is what the steady-state data path allocates. + * tests/alloc_audit/measure_alloc.py is that driver. + * + * The signal handler only uses async-signal-safe calls (open, write, close). + */ +#ifndef _GNU_SOURCE +#define _GNU_SOURCE +#endif +#include +#include +#include +#include +#include +#include +#include +#include + +static uint64_t allocations; +static uint64_t frees; +static uint64_t bytes; +static const char *out_path; + +static void +count_alloc(size_t size) +{ + __atomic_add_fetch(&allocations, 1, __ATOMIC_RELAXED); + __atomic_add_fetch(&bytes, (uint64_t)size, __ATOMIC_RELAXED); +} + +static void +count_free(void *ptr) +{ + if (ptr != NULL) { + __atomic_add_fetch(&frees, 1, __ATOMIC_RELAXED); + } +} + +static char * +format_u64(char *end, uint64_t value) +{ + do { + *--end = (char)('0' + value % 10); + value /= 10; + } while (value != 0); + return end; +} + +static void +snapshot(int signo) +{ + (void)signo; + int saved_errno = errno; + if (out_path != NULL) { + uint64_t fields[4] = { + (uint64_t)getpid(), + __atomic_load_n(&allocations, __ATOMIC_RELAXED), + __atomic_load_n(&frees, __ATOMIC_RELAXED), + __atomic_load_n(&bytes, __ATOMIC_RELAXED), + }; + char line[128]; + char *end = line + sizeof(line); + char *p = end; + *--p = '\n'; + for (int i = 3; i >= 0; i--) { + p = format_u64(p, fields[i]); + if (i != 0) { + *--p = ' '; + } + } + int fd = open(out_path, O_WRONLY | O_CREAT | O_APPEND, 0644); + if (fd >= 0) { + ssize_t ignored = write(fd, p, (size_t)(end - p)); + (void)ignored; + close(fd); + } + } + errno = saved_errno; +} + +__attribute__((constructor)) static void +install(void) +{ + out_path = getenv("SS_ALLOC_AUDIT_OUT"); + struct sigaction action; + memset(&action, 0, sizeof(action)); + action.sa_handler = snapshot; + action.sa_flags = SA_RESTART; + sigemptyset(&action.sa_mask); + sigaction(SIGUSR2, &action, NULL); +} + +#if defined(__APPLE__) + +#include + +static void * +audit_malloc(size_t size) +{ + count_alloc(size); + return malloc(size); +} + +static void * +audit_calloc(size_t count, size_t size) +{ + count_alloc(count * size); + return calloc(count, size); +} + +static void * +audit_realloc(void *ptr, size_t size) +{ + count_alloc(size); + return realloc(ptr, size); +} + +static int +audit_posix_memalign(void **out, size_t align, size_t size) +{ + count_alloc(size); + return posix_memalign(out, align, size); +} + +static void +audit_free(void *ptr) +{ + count_free(ptr); + free(ptr); +} + +#define SS_INTERPOSE(replacement, original) \ + __attribute__((used)) static const struct { \ + const void *replacement_fn; \ + const void *original_fn; \ + } interpose_##original __attribute__((section("__DATA,__interpose"))) = { \ + (const void *)(uintptr_t)&replacement, (const void *)(uintptr_t)&original } + +SS_INTERPOSE(audit_malloc, malloc); +SS_INTERPOSE(audit_calloc, calloc); +SS_INTERPOSE(audit_realloc, realloc); +SS_INTERPOSE(audit_posix_memalign, posix_memalign); +SS_INTERPOSE(audit_free, free); + +#elif defined(__GLIBC__) + +extern void *__libc_malloc(size_t); +extern void *__libc_calloc(size_t, size_t); +extern void *__libc_realloc(void *, size_t); +extern void *__libc_memalign(size_t, size_t); +extern void __libc_free(void *); + +void * +malloc(size_t size) +{ + count_alloc(size); + return __libc_malloc(size); +} + +void * +calloc(size_t count, size_t size) +{ + count_alloc(count * size); + return __libc_calloc(count, size); +} + +void * +realloc(void *ptr, size_t size) +{ + count_alloc(size); + return __libc_realloc(ptr, size); +} + +void * +memalign(size_t align, size_t size) +{ + count_alloc(size); + return __libc_memalign(align, size); +} + +void * +aligned_alloc(size_t align, size_t size) +{ + return memalign(align, size); +} + +int +posix_memalign(void **out, size_t align, size_t size) +{ + void *p = memalign(align, size); + if (p == NULL) { + return ENOMEM; + } + *out = p; + return 0; +} + +void +free(void *ptr) +{ + count_free(ptr); + __libc_free(ptr); +} + +#else +#error "alloc_counter supports glibc and macOS only" +#endif diff --git a/tests/alloc_audit/measure_alloc.py b/tests/alloc_audit/measure_alloc.py new file mode 100644 index 00000000..28c634d7 --- /dev/null +++ b/tests/alloc_audit/measure_alloc.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: GPL-3.0-or-later +"""Count heap allocations per proxied TCP connection in ss-local and ss-server. + +Both peers run with the alloc_counter shim preloaded. After a warm-up the +script snapshots each process's allocation counter, relays more connections, +and snapshots again. The difference, divided by the number of connections, +is what one connection costs in the steady state. The no-heap goal for the +C++ core is zero. + +Report-only by default; --max-per-connection N turns it into a check. + +Example: + python3 tests/alloc_audit/measure_alloc.py --bin build/bin \\ + --shim build/lib/libss_alloc_audit.dylib --method aes-256-gcm +""" +import argparse +import base64 +import concurrent.futures +import json +import os +from pathlib import Path +import signal +import sys +import tempfile +import threading +import time + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from interop import TCPHandler, TCPOrigin, free_port, peers, tcp_case # noqa: E402 + + +def read_snapshot(path, previous_lines, timeout=10): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + lines = path.read_text().splitlines() if path.exists() else [] + if len(lines) > previous_lines: + pid, allocations, frees, size = (int(v) for v in lines[-1].split()) + return {"pid": pid, "allocations": allocations, "frees": frees, "bytes": size}, len(lines) + time.sleep(.02) + raise AssertionError(f"no allocation snapshot written to {path}") + + +def snapshot_all(processes, outputs, counts): + results = [] + for index, (process, path) in enumerate(zip(processes, outputs)): + os.kill(process.pid, signal.SIGUSR2) + result, counts[index] = read_snapshot(path, counts[index]) + results.append(result) + return results + + +def relay(local_port, origin_port, connections, concurrency, size): + with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool: + futures = [pool.submit(tcp_case, local_port, origin_port, size) for _ in range(connections)] + for future in futures: + future.result(timeout=60) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--bin", default="build/bin", help="directory with ss-local and ss-server") + parser.add_argument("--shim", required=True, help="path to the alloc_counter shared library") + parser.add_argument("--method", default="aes-256-gcm") + parser.add_argument("--warmup", type=int, default=50, help="connections before the first snapshot") + parser.add_argument("--connections", type=int, default=200, help="connections between snapshots") + 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("--json", action="store_true", help="print the result as JSON") + args = parser.parse_args() + + binary_dir = Path(args.bin).resolve() + shim = str(Path(args.shim).resolve()) + preload = "DYLD_INSERT_LIBRARIES" if sys.platform == "darwin" else "LD_PRELOAD" + 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, \ + 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] + commands = [ + [str(binary_dir / "ss-server"), "-s", "127.0.0.1", "-p", str(server_port), + "-k", password, "-m", args.method], + [str(binary_dir / "ss-local"), "-s", "127.0.0.1", "-p", str(server_port), + "-l", str(local_port), "-k", password, "-m", args.method], + ] + environments = [] + for output in outputs: + env = os.environ.copy() + env[preload] = shim + env["SS_ALLOC_AUDIT_OUT"] = str(output) + environments.append(env) + + # peers() takes one environment for all of its commands, so each + # peer gets its own context and its own output file. + with peers([commands[0]], [server_port], environments[0]) as server, \ + peers([commands[1]], [local_port], environments[1]) as local: + processes = server + local + counts = [0, 0] + relay(local_port, origin.server_address[1], args.warmup, args.concurrency, args.size) + time.sleep(.5) + before = snapshot_all(processes, outputs, counts) + relay(local_port, origin.server_address[1], args.connections, args.concurrency, args.size) + time.sleep(.5) + after = snapshot_all(processes, outputs, counts) + + report = {"method": args.method, "connections": args.connections, "size": args.size, "peers": {}} + failed = False + for name, b, a in zip(names, before, after): + delta = a["allocations"] - b["allocations"] + per = delta / args.connections + report["peers"][name] = { + "allocations": delta, + "frees": a["frees"] - b["frees"], + "bytes": a["bytes"] - b["bytes"], + "allocations_per_connection": round(per, 2), + } + if args.max_per_connection is not None and per > args.max_per_connection: + failed = True + + if args.json: + print(json.dumps(report, indent=2)) + else: + print(f"method {args.method}, {args.connections} connections of {args.size} bytes") + for name, peer in report["peers"].items(): + print(f" {name:9} {peer['allocations_per_connection']:8.2f} allocations/connection " + f"({peer['allocations']} total, {peer['bytes']} bytes)") + if failed: + print(f"FAIL: more than {args.max_per_connection} allocations per connection") + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/interop.py b/tests/interop.py index c9a04083..7b7efb48 100644 --- a/tests/interop.py +++ b/tests/interop.py @@ -159,7 +159,7 @@ def peers(commands, ports, environment=None): process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=environment) processes.append(process) wait_ready(process, port) - yield + yield processes except BaseException: for path in logs: print(path.name + ":\n" + path.read_text(errors="replace")[-6000:], flush=True)