commit 40b892886bdfbbdcf8dab561e142925f1f27dee9
| author | 斟酌 鵬兄 <tgckpg@gmail.com> |
|---|
| date | 2026-05-26T14:56:37Z |
|---|
| subject | Fixed and added tests for proxy v2 |
|---|
commit 40b892886bdfbbdcf8dab561e142925f1f27dee9
Author: 斟酌 鵬兄 <tgckpg@gmail.com>
Date: 2026-05-26T14:56:37Z
Fixed and added tests for proxy v2
---
.github/workflows/cmake-multi-platform.yml | 2 +-
.gitignore | 1 +
Dockerfile | 3 +
Makefile | 2 +-
file_conf.c | 2 +-
route.c | 2 +-
tcp_route.c | 19 +-
tests/__init__.py | 0
tests/run_tests.py | 76 ++++++++
tests/support.py | 264 ++++++++++++++++++++++++++++
tests/test_haproxy_proxy_v2.py | 116 +++++++++++++
tests/test_proxy.py | 268 -----------------------------
tests/test_tcp_basic.py | 30 ++++
tests/test_tcp_stress.py | 48 ++++++
tinyproxy.conf | 4 +-
15 files changed, 555 insertions(+), 282 deletions(-)
diff --git a/.github/workflows/cmake-multi-platform.yml b/.github/workflows/cmake-multi-platform.yml
index 5093d38..aa052ca 100644
--- a/.github/workflows/cmake-multi-platform.yml
+++ b/.github/workflows/cmake-multi-platform.yml
@@ -94,7 +94,7 @@ jobs:
- name: Install deps
run: |
brew update
- brew install libevent
+ brew install libevent haproxy
- name: Build
run: make clean all STATIC=0
diff --git a/.gitignore b/.gitignore
index 95d0a5f..63ce10f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,6 +1,7 @@
bin/
build/
libevent2/
+__pycache__
*.log
*.swp
diff --git a/Dockerfile b/Dockerfile
index 9d7b853..0fcf32c 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -10,6 +10,9 @@ COPY [ "*.c", "*.h", "*.conf", "Makefile", "/src" ]
RUN make all STATIC=1
+RUN --mount=type=cache,target=/var/cache/apk,sharing=locked \
+ apk add --update-cache haproxy
+
COPY tests ./tests
RUN make test STATIC=1
diff --git a/Makefile b/Makefile
index cbec336..183261d 100644
--- a/Makefile
+++ b/Makefile
@@ -107,7 +107,7 @@ else
endif
test: $(BIN)
- $(TEST_FLAGS) python3 tests/test_proxy.py $(BIN)
+ $(TEST_FLAGS) python3 -m tests.run_tests $(BIN)
clean:
rm -rf $(BIN_DIR)
diff --git a/file_conf.c b/file_conf.c
index 11d125c..50a52b8 100644
--- a/file_conf.c
+++ b/file_conf.c
@@ -130,7 +130,7 @@ static int parse_route_line(char *line, struct route *route)
route->send_proxy_v2 = false;
if (count == 4) {
- if (strcmp(fields[3], "proxyv2") != 0) {
+ if (strcmp(fields[3], "proxy_v2") != 0) {
return -1;
}
diff --git a/route.c b/route.c
index 61e4238..eb39569 100644
--- a/route.c
+++ b/route.c
@@ -16,7 +16,7 @@ void route_options_str(const struct route *r, char *buf, size_t buflen)
} while (0)
if (r->send_proxy_v2) {
- ADD_OPT("proxyv2");
+ ADD_OPT("proxy_v2");
}
#undef ADD_OPT
diff --git a/tcp_route.c b/tcp_route.c
index d0d1cc6..554adba 100644
--- a/tcp_route.c
+++ b/tcp_route.c
@@ -16,8 +16,6 @@
#include "route.h"
#include "tcp_route.h"
-#define PROXY_V2_SIG "\r\n\r\n\0\r\nQUIT\n"
-
static void free_conn(conn_t *conn) {
if (conn == NULL) {
return;
@@ -154,7 +152,7 @@ static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac)
return;
}
- if (0 < ac->peer_addr_len) {
+ if (ac->peer_addr_len <= 0) {
LOG_ERROR("invalid client address");
free_conn(conn);
return;
@@ -167,15 +165,17 @@ static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac)
return;
}
- if (proxy_v2_write_bufferevent(
+ int rc = proxy_v2_write_bufferevent(
conn->upstream,
(const struct sockaddr *)&ac->peer_addr,
ac->peer_addr_len,
(struct sockaddr *)&local_addr,
local_len,
SOCK_STREAM
- ) < 0) {
- LOG_ERROR("failed to write PROXY v2 header");
+ );
+
+ if (rc < 0) {
+ LOG_ERROR("failed to write PROXY v2 header", "err", _LOGV(strerror(-rc)));
free_conn(conn);
return;
}
@@ -204,10 +204,13 @@ static void accept_cb(
.route = ctx->route,
};
- if(addr && 0 < socklen && (size_t)socklen <= sizeof(ac.peer_addr)) {
+ if (addr != NULL && socklen > 0 && (size_t)socklen <= sizeof(ac.peer_addr)) {
memcpy(&ac.peer_addr, addr, (size_t)socklen);
+ ac.peer_addr_len = (socklen_t)socklen;
} else {
- ac.peer_addr_len = 0;
+ LOG_ERROR("invalid accepted client address", "socklen", _LOGV(socklen));
+ evutil_closesocket(client_fd);
+ return;
}
dispatch_client_fd(ctx->worker, &ac);
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/run_tests.py b/tests/run_tests.py
new file mode 100644
index 0000000..4d5d170
--- /dev/null
+++ b/tests/run_tests.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+
+import asyncio
+import os
+import sys
+
+from .support import SkipTest, raise_fd_limit, run_default_tinyproxy
+from . import test_tcp_basic
+from . import test_tcp_stress
+from . import test_haproxy_proxy_v2
+
+
+DEFAULT_PROXY_TEST_MODULES = [
+ test_tcp_basic,
+ test_tcp_stress,
+]
+
+STANDALONE_TEST_MODULES = [
+ test_haproxy_proxy_v2,
+]
+
+
+async def run_module_tests(module) -> tuple[int, int]:
+ passed = 0
+ skipped = 0
+
+ for name, test_func in module.TESTS:
+ print(f"running {name}...")
+
+ try:
+ await test_func()
+ except SkipTest as exc:
+ skipped += 1
+ print(f"skipped {name}: {exc}")
+ continue
+
+ passed += 1
+ print(f"ok {name}")
+
+ return passed, skipped
+
+
+async def main_async(proxy_bin: str) -> int:
+ os.environ["TINYPROXY_BIN"] = proxy_bin
+
+ fd_limit = int(os.environ.get("FD_LIMIT", "65535"))
+ raise_fd_limit(fd_limit)
+
+ passed = 0
+ skipped = 0
+
+ async with run_default_tinyproxy(proxy_bin):
+ for module in DEFAULT_PROXY_TEST_MODULES:
+ p, s = await run_module_tests(module)
+ passed += p
+ skipped += s
+
+ for module in STANDALONE_TEST_MODULES:
+ p, s = await run_module_tests(module)
+ passed += p
+ skipped += s
+
+ print(f"all tests passed: passed={passed} skipped={skipped}")
+ return 0
+
+
+def main() -> int:
+ if len(sys.argv) != 2:
+ print(f"usage: {sys.argv[0]} tinyproxy", file=sys.stderr)
+ return 2
+
+ return asyncio.run(main_async(sys.argv[1]))
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/tests/support.py b/tests/support.py
new file mode 100644
index 0000000..d1d1c37
--- /dev/null
+++ b/tests/support.py
@@ -0,0 +1,264 @@
+import asyncio
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+
+try:
+ import resource
+except ImportError:
+ resource = None
+
+
+LISTEN_HOST = "127.0.0.1"
+PROXY_PORT = 31232
+BACKEND_PORT = 41232
+
+
+class SkipTest(Exception):
+ pass
+
+
+def raise_fd_limit(wanted: int) -> None:
+ if resource is None:
+ return
+
+ soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+
+ if soft >= wanted:
+ return
+
+ new_soft = min(wanted, hard)
+
+ try:
+ resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
+ print(f"raised RLIMIT_NOFILE soft limit: {soft} -> {new_soft}")
+ except OSError as exc:
+ print(f"warning: failed to raise fd limit: {exc}", file=sys.stderr)
+ print(f"current RLIMIT_NOFILE soft={soft} hard={hard}", file=sys.stderr)
+
+
+def find_program(env_name: str, fallback_name: str) -> str | None:
+ from_env = os.environ.get(env_name)
+ if from_env:
+ return from_env
+
+ return shutil.which(fallback_name)
+
+
+async def wait_for_port(host: str, port: int, timeout: float = 5.0) -> None:
+ deadline = time.monotonic() + timeout
+
+ while time.monotonic() < deadline:
+ try:
+ reader, writer = await asyncio.open_connection(host, port)
+ writer.close()
+ await writer.wait_closed()
+ return
+ except OSError:
+ await asyncio.sleep(0.05)
+
+ raise RuntimeError(f"timed out waiting for {host}:{port}")
+
+
+async def echo_handler(
+ reader: asyncio.StreamReader,
+ writer: asyncio.StreamWriter,
+) -> None:
+ try:
+ while True:
+ data = await reader.read(65536)
+ if not data:
+ return
+
+ writer.write(data)
+ await writer.drain()
+
+ except (ConnectionResetError, BrokenPipeError):
+ return
+
+ finally:
+ writer.close()
+
+ try:
+ await writer.wait_closed()
+ except OSError:
+ pass
+
+
+async def recv_exact(reader: asyncio.StreamReader, size: int) -> bytes:
+ chunks = []
+ remaining = size
+
+ while remaining > 0:
+ chunk = await reader.read(min(65536, remaining))
+ if not chunk:
+ break
+
+ chunks.append(chunk)
+ remaining -= len(chunk)
+
+ return b"".join(chunks)
+
+
+async def proxy_roundtrip(
+ payload: bytes,
+ host: str = LISTEN_HOST,
+ port: int = PROXY_PORT,
+ timeout: float = 10.0,
+) -> bytes:
+ reader, writer = await asyncio.open_connection(host, port)
+
+ try:
+ writer.write(payload)
+ await writer.drain()
+
+ return await asyncio.wait_for(
+ recv_exact(reader, len(payload)),
+ timeout=timeout,
+ )
+
+ finally:
+ writer.close()
+
+ try:
+ await writer.wait_closed()
+ except OSError:
+ pass
+
+
+def write_temp_file(content: str, suffix: str) -> str:
+ fd, path = tempfile.mkstemp(suffix=suffix, text=True)
+
+ with os.fdopen(fd, "w") as f:
+ f.write(content)
+
+ return path
+
+
+def terminate_process(name: str, proc: subprocess.Popen | None) -> None:
+ if proc is None:
+ return
+
+ if proc.poll() is not None:
+ return
+
+ proc.terminate()
+
+ try:
+ proc.wait(timeout=2.0)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ proc.wait(timeout=2.0)
+
+
+def print_process_output(name: str, proc: subprocess.Popen | None) -> None:
+ if proc is None:
+ return
+
+ stdout = proc.stdout.read() if proc.stdout else ""
+ stderr = proc.stderr.read() if proc.stderr else ""
+
+ if stdout:
+ print(f"\n{name} stdout:")
+ print(stdout)
+
+ if stderr:
+ print(f"\n{name} stderr:")
+ print(stderr)
+
+
+@dataclass
+class TinyproxyFixture:
+ proxy: subprocess.Popen
+ conf_path: str
+ listen_host: str
+ listen_port: int
+
+
+@asynccontextmanager
+async def run_echo_backend(
+ host: str = LISTEN_HOST,
+ port: int = BACKEND_PORT,
+):
+ server = await asyncio.start_server(
+ echo_handler,
+ host,
+ port,
+ backlog=4096,
+ )
+
+ try:
+ await wait_for_port(host, port)
+ yield server
+ finally:
+ server.close()
+ await server.wait_closed()
+
+
+@asynccontextmanager
+async def run_tinyproxy_with_conf(
+ proxy_bin: str,
+ conf_text: str,
+ listen_host: str = LISTEN_HOST,
+ listen_port: int = PROXY_PORT,
+):
+ conf_path = None
+ proxy = None
+
+ try:
+ conf_path = write_temp_file(conf_text, ".conf")
+
+ proxy = subprocess.Popen(
+ [proxy_bin, "-c", conf_path],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ )
+
+ await wait_for_port(listen_host, listen_port)
+
+ if proxy.poll() is not None:
+ stderr = proxy.stderr.read() if proxy.stderr else ""
+ raise RuntimeError(
+ f"tinyproxy exited early with code {proxy.returncode}\n{stderr}"
+ )
+
+ yield TinyproxyFixture(
+ proxy=proxy,
+ conf_path=conf_path,
+ listen_host=listen_host,
+ listen_port=listen_port,
+ )
+
+ finally:
+ terminate_process("tinyproxy", proxy)
+ print_process_output("tinyproxy", proxy)
+
+ if conf_path is not None:
+ try:
+ os.unlink(conf_path)
+ except FileNotFoundError:
+ pass
+
+
+@asynccontextmanager
+async def run_default_tinyproxy(proxy_bin: str):
+ conf_text = (
+ f"{LISTEN_HOST}:{PROXY_PORT} "
+ f"{LISTEN_HOST}:{BACKEND_PORT} "
+ f"tcp\n"
+ )
+
+ async with run_echo_backend(LISTEN_HOST, BACKEND_PORT):
+ async with run_tinyproxy_with_conf(
+ proxy_bin=proxy_bin,
+ conf_text=conf_text,
+ listen_host=LISTEN_HOST,
+ listen_port=PROXY_PORT,
+ ) as fixture:
+ yield fixture
diff --git a/tests/test_haproxy_proxy_v2.py b/tests/test_haproxy_proxy_v2.py
new file mode 100644
index 0000000..fab0dfa
--- /dev/null
+++ b/tests/test_haproxy_proxy_v2.py
@@ -0,0 +1,116 @@
+import os
+import sys
+import subprocess
+
+from .support import (
+ BACKEND_PORT,
+ LISTEN_HOST,
+ PROXY_PORT,
+ SkipTest,
+ find_program,
+ print_process_output,
+ proxy_roundtrip,
+ run_echo_backend,
+ run_tinyproxy_with_conf,
+ terminate_process,
+ wait_for_port,
+ write_temp_file,
+)
+
+
+HAPROXY_FRONTEND_PORT = 41233
+
+
+def require_haproxy() -> str:
+ if sys.platform.startswith("win"):
+ raise SkipTest("HAProxy proxy-v2 tests are skipped on Windows")
+
+ haproxy_bin = find_program("HAPROXY_BIN", "haproxy")
+ if haproxy_bin is None:
+ raise SkipTest("HAProxy not found; set HAPROXY_BIN to enable this test")
+
+ return haproxy_bin
+
+
+async def test_tinyproxy_sends_proxy_v2_to_haproxy() -> None:
+ haproxy_bin = require_haproxy()
+
+ haproxy_conf = f"""
+global
+ maxconn 4096
+
+defaults
+ mode tcp
+ timeout connect 5s
+ timeout client 30s
+ timeout server 30s
+
+frontend proxy_v2_in
+ bind {LISTEN_HOST}:{HAPROXY_FRONTEND_PORT} accept-proxy
+ default_backend echo_backend
+
+backend echo_backend
+ server echo1 {LISTEN_HOST}:{BACKEND_PORT}
+"""
+
+ haproxy_conf_path = None
+ haproxy = None
+
+ # Adjust this if your config parser uses a different spelling.
+ #
+ # Current expected route shape:
+ #
+ # listen upstream tcp proxy_v2
+ #
+ proxy_v2_option = os.environ.get("TINYPROXY_PROXY_V2_OPTION", "proxy_v2")
+
+ tinyproxy_conf = (
+ f"{LISTEN_HOST}:{PROXY_PORT} "
+ f"{LISTEN_HOST}:{HAPROXY_FRONTEND_PORT} "
+ f"tcp {proxy_v2_option}\n"
+ )
+
+ try:
+ haproxy_conf_path = write_temp_file(haproxy_conf, ".haproxy.cfg")
+
+ haproxy = subprocess.Popen(
+ [haproxy_bin, "-f", haproxy_conf_path, "-db"],
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ text=True,
+ )
+
+ async with run_echo_backend(LISTEN_HOST, BACKEND_PORT):
+ await wait_for_port(LISTEN_HOST, HAPROXY_FRONTEND_PORT)
+
+ if haproxy.poll() is not None:
+ stderr = haproxy.stderr.read() if haproxy.stderr else ""
+ raise RuntimeError(
+ f"haproxy exited early with code {haproxy.returncode}\n{stderr}"
+ )
+
+ async with run_tinyproxy_with_conf(
+ proxy_bin=os.environ["TINYPROXY_BIN"],
+ conf_text=tinyproxy_conf,
+ listen_host=LISTEN_HOST,
+ listen_port=PROXY_PORT,
+ ):
+ payload = b"hello through proxy v2 via haproxy\n"
+ got = await proxy_roundtrip(payload, port=PROXY_PORT)
+
+ assert got == payload, f"proxy-v2 haproxy roundtrip mismatch: {got!r}"
+
+ finally:
+ terminate_process("haproxy", haproxy)
+ print_process_output("haproxy", haproxy)
+
+ if haproxy_conf_path is not None:
+ try:
+ os.unlink(haproxy_conf_path)
+ except FileNotFoundError:
+ pass
+
+
+TESTS = [
+ ("test_tinyproxy_sends_proxy_v2_to_haproxy", test_tinyproxy_sends_proxy_v2_to_haproxy),
+]
diff --git a/tests/test_proxy.py b/tests/test_proxy.py
deleted file mode 100755
index f48bdb8..0000000
--- a/tests/test_proxy.py
+++ /dev/null
@@ -1,268 +0,0 @@
-#!/usr/bin/env python3
-
-import asyncio
-import os
-import signal
-import subprocess
-import sys
-import time
-import tempfile
-
-try:
- import resource
-except ImportError:
- resource = None
-
-LISTEN_HOST = "127.0.0.1"
-PROXY_PORT = 31232
-BACKEND_PORT = 41232
-
-def raise_fd_limit(wanted: int) -> None:
- if resource is None:
- return
-
- soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
-
- if soft >= wanted:
- return
-
- new_soft = min(wanted, hard)
-
- try:
- resource.setrlimit(resource.RLIMIT_NOFILE, (new_soft, hard))
- print(f"raised RLIMIT_NOFILE soft limit: {soft} -> {new_soft}")
- except OSError as exc:
- print(f"warning: failed to raise fd limit: {exc}", file=sys.stderr)
- print(f"current RLIMIT_NOFILE soft={soft} hard={hard}", file=sys.stderr)
-
-
-async def wait_for_port(host: str, port: int, timeout: float = 5.0) -> None:
- deadline = time.monotonic() + timeout
-
- while time.monotonic() < deadline:
- try:
- reader, writer = await asyncio.open_connection(host, port)
- writer.close()
- await writer.wait_closed()
- return
- except OSError:
- await asyncio.sleep(0.05)
-
- raise RuntimeError(f"timed out waiting for {host}:{port}")
-
-
-async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
- try:
- while True:
- data = await reader.read(65536)
- if not data:
- return
-
- writer.write(data)
- await writer.drain()
-
- except (ConnectionResetError, BrokenPipeError):
- return
-
- finally:
- writer.close()
-
- try:
- await writer.wait_closed()
- except OSError:
- pass
-
-
-async def recv_exact(reader: asyncio.StreamReader, size: int) -> bytes:
- chunks = []
- remaining = size
-
- while remaining > 0:
- chunk = await reader.read(min(65536, remaining))
- if not chunk:
- break
-
- chunks.append(chunk)
- remaining -= len(chunk)
-
- return b"".join(chunks)
-
-
-async def proxy_roundtrip(payload: bytes, timeout: float = 10.0) -> bytes:
- reader, writer = await asyncio.open_connection(LISTEN_HOST, PROXY_PORT)
-
- try:
- writer.write(payload)
- await writer.drain()
-
- return await asyncio.wait_for(
- recv_exact(reader, len(payload)),
- timeout=timeout,
- )
-
- finally:
- writer.close()
-
- try:
- await writer.wait_closed()
- except OSError:
- pass
-
-
-async def test_small_roundtrip() -> None:
- payload = b"hello through proxy\n"
- got = await proxy_roundtrip(payload)
- assert got == payload, f"small roundtrip mismatch: {got!r}"
-
-
-async def test_large_roundtrip() -> None:
- payload = b"0123456789abcdef" * 131072 # 2 MiB
- got = await proxy_roundtrip(payload, timeout=30.0)
- assert got == payload, f"large roundtrip mismatch: got {len(got)} bytes"
-
-
-async def test_many_sequential_connections(count: int) -> None:
- for i in range(count):
- payload = f"message-{i}\n".encode()
- got = await proxy_roundtrip(payload)
- assert got == payload, f"sequential connection {i} failed"
-
-
-async def run_one_concurrent(i: int, payload_size: int, sem: asyncio.Semaphore) -> None:
- async with sem:
- prefix = f"worker-{i}-".encode()
- payload = (prefix * ((payload_size // len(prefix)) + 1))[:payload_size]
-
- got = await proxy_roundtrip(payload, timeout=30.0)
-
- if got != payload:
- raise AssertionError(
- f"worker {i} mismatch: expected {len(payload)} bytes, got {len(got)}"
- )
-
-
-async def test_concurrent_connections(
- total: int,
- concurrency: int,
- payload_size: int,
-) -> None:
- sem = asyncio.Semaphore(concurrency)
-
- tasks = [
- asyncio.create_task(run_one_concurrent(i, payload_size, sem))
- for i in range(total)
- ]
-
- done = 0
-
- for task in asyncio.as_completed(tasks):
- await task
- done += 1
-
- if done % 1000 == 0:
- print(f" completed {done}/{total}")
-
-
-async def main_async(proxy_bin: str) -> int:
- fd_limit = int(os.environ.get("FD_LIMIT", "65535"))
- raise_fd_limit(fd_limit)
-
- backend_server = await asyncio.start_server(
- echo_handler,
- LISTEN_HOST,
- BACKEND_PORT,
- backlog=4096,
- )
-
- conf_path = None
- proxy = None
-
- try:
- fd, conf_path = tempfile.mkstemp(suffix=".conf", text=True)
- with os.fdopen(fd, "w") as f:
- f.write(f"{LISTEN_HOST}:{PROXY_PORT} {LISTEN_HOST}:{BACKEND_PORT} tcp\n")
-
- proxy = subprocess.Popen(
- [proxy_bin, "-c", conf_path],
- stdout=subprocess.PIPE,
- stderr=subprocess.PIPE,
- text=True,
- )
-
- await asyncio.sleep(0.2)
-
- if proxy.poll() is not None:
- stderr = proxy.stderr.read() if proxy.stderr else ""
- raise RuntimeError(
- f"proxy exited early with code {proxy.returncode}\n{stderr}"
- )
- await wait_for_port(LISTEN_HOST, BACKEND_PORT)
- await wait_for_port(LISTEN_HOST, PROXY_PORT)
-
- print("running test_small_roundtrip...")
- await test_small_roundtrip()
- print("ok test_small_roundtrip")
-
- print("running test_large_roundtrip...")
- await test_large_roundtrip()
- print("ok test_large_roundtrip")
-
- print("running test_many_sequential_connections 1000...")
- await test_many_sequential_connections(1000)
- print("ok test_many_sequential_connections")
-
- total = int(os.environ.get("TOTAL", "10000"))
- concurrency = int(os.environ.get("CONCURRENCY", "10000"))
- payload_size = int(os.environ.get("PAYLOAD_SIZE", "1024"))
-
- print(
- f"running test_concurrent_connections "
- f"total={total} concurrency={concurrency} payload_size={payload_size}..."
- )
-
- await test_concurrent_connections(
- total=total,
- concurrency=concurrency,
- payload_size=payload_size,
- )
-
- print("ok test_concurrent_connections")
- print("all tests passed")
- return 0
-
- finally:
- if proxy is not None:
- proxy.terminate()
-
- try:
- proxy.wait(timeout=2.0)
- except subprocess.TimeoutExpired:
- proxy.kill()
- proxy.wait(timeout=2.0)
-
- backend_server.close()
- await backend_server.wait_closed()
-
- if proxy is not None:
- stderr = proxy.stderr.read() if proxy.stderr else ""
- if stderr:
- print("\nproxy stderr:")
- print(stderr)
-
- if conf_path is not None:
- try:
- os.unlink(conf_path)
- except FileNotFoundError:
- pass
-
-
-def main() -> int:
- if len(sys.argv) != 2:
- print(f"usage: {sys.argv[0]} tinyproxy", file=sys.stderr)
- return 2
-
- return asyncio.run(main_async(sys.argv[1]))
-
-
-if __name__ == "__main__":
- raise SystemExit(main())
diff --git a/tests/test_tcp_basic.py b/tests/test_tcp_basic.py
new file mode 100644
index 0000000..c928391
--- /dev/null
+++ b/tests/test_tcp_basic.py
@@ -0,0 +1,30 @@
+from .support import proxy_roundtrip
+
+
+async def test_small_roundtrip() -> None:
+ payload = b"hello through proxy\n"
+ got = await proxy_roundtrip(payload)
+
+ assert got == payload, f"small roundtrip mismatch: {got!r}"
+
+
+async def test_large_roundtrip() -> None:
+ payload = b"0123456789abcdef" * 131072 # 2 MiB
+ got = await proxy_roundtrip(payload, timeout=30.0)
+
+ assert got == payload, f"large roundtrip mismatch: got {len(got)} bytes"
+
+
+async def test_many_sequential_connections(count: int = 1000) -> None:
+ for i in range(count):
+ payload = f"message-{i}\n".encode()
+ got = await proxy_roundtrip(payload)
+
+ assert got == payload, f"sequential connection {i} failed"
+
+
+TESTS = [
+ ("test_small_roundtrip", test_small_roundtrip),
+ ("test_large_roundtrip", test_large_roundtrip),
+ ("test_many_sequential_connections", test_many_sequential_connections),
+]
diff --git a/tests/test_tcp_stress.py b/tests/test_tcp_stress.py
new file mode 100644
index 0000000..456b260
--- /dev/null
+++ b/tests/test_tcp_stress.py
@@ -0,0 +1,48 @@
+import asyncio
+import os
+
+from .support import proxy_roundtrip
+
+
+async def run_one_concurrent(
+ i: int,
+ payload_size: int,
+ sem: asyncio.Semaphore,
+) -> None:
+ async with sem:
+ prefix = f"worker-{i}-".encode()
+ payload = (prefix * ((payload_size // len(prefix)) + 1))[:payload_size]
+
+ got = await proxy_roundtrip(payload, timeout=30.0)
+
+ if got != payload:
+ raise AssertionError(
+ f"worker {i} mismatch: expected {len(payload)} bytes, got {len(got)}"
+ )
+
+
+async def test_concurrent_connections() -> None:
+ total = int(os.environ.get("TOTAL", "10000"))
+ concurrency = int(os.environ.get("CONCURRENCY", "10000"))
+ payload_size = int(os.environ.get("PAYLOAD_SIZE", "1024"))
+
+ sem = asyncio.Semaphore(concurrency)
+
+ tasks = [
+ asyncio.create_task(run_one_concurrent(i, payload_size, sem))
+ for i in range(total)
+ ]
+
+ done = 0
+
+ for task in asyncio.as_completed(tasks):
+ await task
+ done += 1
+
+ if done % 1000 == 0:
+ print(f" completed {done}/{total}")
+
+
+TESTS = [
+ ("test_concurrent_connections", test_concurrent_connections),
+]
diff --git a/tinyproxy.conf b/tinyproxy.conf
index df3f5e0..c7e7307 100644
--- a/tinyproxy.conf
+++ b/tinyproxy.conf
@@ -1,5 +1,5 @@
# listen upstream proto option
:80 10.0.1.1:80 tcp
-:443 10.0.1.1:443 tcp proxyv2
+:443 10.0.1.1:443 tcp proxy_v2
:12345 10.0.1.1:12345 udp
-:19132 10.0.1.1:19132 udp proxyv2
+:19132 10.0.1.1:19132 udp proxy_v2