commit 1b97c6dd56b770f3497ce5f60a185ea52e88a1bf
| author | 斟酌 鵬兄 <tgckpg@gmail.com> |
|---|
| date | 2026-05-29T13:41:36Z |
|---|
| subject | Split into stream_pipe.c |
|---|
commit 1b97c6dd56b770f3497ce5f60a185ea52e88a1bf
Author: 斟酌 鵬兄 <tgckpg@gmail.com>
Date: 2026-05-29T13:41:36Z
Split into stream_pipe.c
---
src/klog.h | 8 ++
src/stream_conn.c | 148 +--------------------------------
src/stream_pipe.c | 166 ++++++++++++++++++++++++++++++++++++++
src/stream_pipe.h | 11 +++
tests/support.py | 91 +++++++++++++++++++--
tests/test_tcp_backpressure.py | 11 ++-
tests/test_tcp_connect_timeout.py | 6 +-
tests/test_tcp_idle_timeout.py | 11 ++-
tests/test_tcp_keep_alive.py | 11 ++-
tests/test_unix_listeners.py | 21 +----
10 files changed, 294 insertions(+), 190 deletions(-)
diff --git a/src/klog.h b/src/klog.h
index 1920cb0..5601a5f 100644
--- a/src/klog.h
+++ b/src/klog.h
@@ -105,4 +105,12 @@ static inline struct log_value log_value_endpoint(const struct endpoint *v)
#define LOG_WARN(msg, ...) log_at('W', __FILE__, __LINE__, (msg), ##__VA_ARGS__, NULL)
#define LOG_ERROR(msg, ...) log_at('E', __FILE__, __LINE__, (msg), ##__VA_ARGS__, NULL)
+#ifdef TINYPROXY_DEBUG
+#define LOG_DEBUG(msg, ...) \
+ LOG_INFO((msg), __VA_ARGS__)
+#else
+#define LOG_DEBUG(msg, ...) \
+ do { } while (0)
+#endif
+
#endif
diff --git a/src/stream_conn.c b/src/stream_conn.c
index 5c3701a..6c89588 100644
--- a/src/stream_conn.c
+++ b/src/stream_conn.c
@@ -10,6 +10,7 @@
#include "stream_conn.h"
#include "stream_file.h"
#include "stream_builtin.h"
+#include "stream_pipe.h"
#include "proxy_proto_v2.h"
void free_conn(conn_t *conn) {
@@ -70,19 +71,6 @@ static int set_socket_keepalive(evutil_socket_t fd, const struct route *r)
return 0;
}
-#ifdef TINYPROXY_DEBUG
-static const char *bev_side(conn_t *conn, struct bufferevent *bev)
-{
- if (bev == conn->client) {
- return "client";
- }
- if (bev == conn->upstream) {
- return "upstream";
- }
- return "unknown";
-}
-#endif
-
static size_t bev_output_len(struct bufferevent *bev)
{
if (bev == NULL) {
@@ -238,136 +226,6 @@ static int connect_upstream(struct bufferevent *bev, const struct endpoint *ep)
}
}
-static void pipe_read_cb(struct bufferevent *src, void *arg)
-{
- conn_t *conn = arg;
- struct bufferevent *dst;
-
- if (src == conn->client) {
- dst = conn->upstream;
- } else if (src == conn->upstream) {
- dst = conn->client;
- } else {
- return;
- }
-
- struct evbuffer *input = bufferevent_get_input(src);
- struct evbuffer *output = bufferevent_get_output(dst);
-
-#ifdef TINYPROXY_DEBUG
- const struct route *r = conn->route;
-
- size_t input_len = evbuffer_get_length(input);
- size_t output_before = evbuffer_get_length(output);
-
- LOG_INFO("stream pipe read",
- "line", _LOGV(r->line_no),
- "from", _LOGV(bev_side(conn, src)),
- "to", _LOGV(bev_side(conn, dst)),
- "input_len", _LOGV(input_len),
- "dst_output_before", _LOGV(output_before)
- );
-#endif
-
- evbuffer_add_buffer(output, input);
-
- size_t output_after = evbuffer_get_length(output);
-
-#ifdef TINYPROXY_DEBUG
- LOG_INFO("stream pipe queued",
- "line", _LOGV(r->line_no),
- "from", _LOGV(bev_side(conn, src)),
- "to", _LOGV(bev_side(conn, dst)),
- "dst_output_after", _LOGV(output_after)
- );
-#endif
-
- if (output_after >= BEV_READ_HIGH_WATER) {
-#ifdef TINYPROXY_DEBUG
- LOG_INFO("stream pipe backpressure pause",
- "line", _LOGV(r->line_no),
- "paused", _LOGV(bev_side(conn, src)),
- "dst_output_len", _LOGV(output_after)
- );
-#endif
-
- bufferevent_disable(src, EV_READ);
- }
-}
-
-static void finish_client_write(conn_t *conn)
-{
- evutil_socket_t fd = bufferevent_getfd(conn->client);
-
- if (fd >= 0) {
-#ifndef _WIN32
- shutdown(fd, SHUT_WR);
-#else
- shutdown(fd, SD_SEND);
-#endif
- }
-
- bufferevent_disable(conn->client, EV_WRITE);
-
- /*
- * Keep EV_READ enabled so we can observe client EOF instead of
- * closing with unread data and causing RST on some platforms.
- */
- bufferevent_enable(conn->client, EV_READ);
-}
-
-static void pipe_write_cb(struct bufferevent *dst, void *arg)
-{
- conn_t *conn = arg;
- struct bufferevent *src;
-
- if (dst == conn->client) {
- src = conn->upstream;
- } else if (dst == conn->upstream) {
- src = conn->client;
- } else {
- return;
- }
-
- struct evbuffer *output = bufferevent_get_output(dst);
- size_t output_len = evbuffer_get_length(output);
-
-#ifdef TINYPROXY_DEBUG
- LOG_INFO("stream pipe write",
- "line", _LOGV(conn->route->line_no),
- "dst", _LOGV(bev_side(conn, dst)),
- "src", _LOGV(bev_side(conn, src)),
- "dst_output_len", _LOGV(output_len)
- );
-#endif
-
- if (dst == conn->client &&
- conn->close_client_after_drain &&
- output_len == 0) {
- LOG_INFO("client output drained; shutting down write side",
- "line", _LOGV(conn->route->line_no)
- );
-
- finish_client_write(conn);
- conn->close_client_after_drain = false;
- conn->close_after_client_eof = true;
- return;
- }
-
- if (src != NULL && output_len < BEV_WRITE_RESUME_WATER) {
-#ifdef TINYPROXY_DEBUG
- LOG_INFO("stream pipe backpressure resume",
- "line", _LOGV(conn->route->line_no),
- "resumed", _LOGV(bev_side(conn, src)),
- "dst", _LOGV(bev_side(conn, dst)),
- "dst_output_len", _LOGV(output_len)
- );
-#endif
-
- bufferevent_enable(src, EV_READ);
- }
-}
-
static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac) {
conn_t *conn = calloc(1, sizeof(*conn));
if (conn == NULL) {
@@ -416,8 +274,8 @@ static void worker_adopt_client_fd(struct worker *w, struct accepted_client *ac)
bufferevent_setwatermark(conn->client, EV_READ, 0, BEV_READ_HIGH_WATER);
bufferevent_setwatermark(conn->upstream, EV_READ, 0, BEV_READ_HIGH_WATER);
- bufferevent_setcb(conn->client, pipe_read_cb, pipe_write_cb, event_cb, conn);
- bufferevent_setcb(conn->upstream, pipe_read_cb, pipe_write_cb, event_cb, conn);
+ bufferevent_setcb(conn->client, pipe_client_read_cb, pipe_client_write_cb, event_cb, conn);
+ bufferevent_setcb(conn->upstream, pipe_upstream_read_cb, pipe_upstream_write_cb, event_cb, conn);
/*
* Do not read from the client yet.
diff --git a/src/stream_pipe.c b/src/stream_pipe.c
new file mode 100644
index 0000000..7f4b5f5
--- /dev/null
+++ b/src/stream_pipe.c
@@ -0,0 +1,166 @@
+#include <event2/buffer.h>
+
+#include "klog.h"
+#include "route.h"
+#include "stream_conn.h"
+#include "stream_pipe.h"
+
+static void finish_client_write(conn_t *conn)
+{
+ evutil_socket_t fd = bufferevent_getfd(conn->client);
+
+ if (fd >= 0) {
+#ifndef _WIN32
+ shutdown(fd, SHUT_WR);
+#else
+ shutdown(fd, SD_SEND);
+#endif
+ }
+
+ bufferevent_disable(conn->client, EV_WRITE);
+
+ /*
+ * Keep EV_READ enabled so we can observe client EOF instead of
+ * closing with unread data and causing RST on some platforms.
+ */
+ bufferevent_enable(conn->client, EV_READ);
+}
+
+void pipe_client_read_cb(struct bufferevent *client, void *arg)
+{
+ conn_t *conn = arg;
+ struct bufferevent *upstream = conn->upstream;
+
+ if (upstream == NULL) {
+ return;
+ }
+
+ struct evbuffer *src = bufferevent_get_input(client);
+ struct evbuffer *dst = bufferevent_get_output(upstream);
+
+#ifdef TINYPROXY_DEBUG
+ size_t len = evbuffer_get_length(src);
+
+ LOG_DEBUG("stream pipe read",
+ "line", _LOGV(conn->route->line_no),
+ "src", _LOGV("client"),
+ "dst", _LOGV("upstream"),
+ "bytes", _LOGV(len)
+ );
+#endif
+
+ evbuffer_add_buffer(dst, src);
+
+ if (evbuffer_get_length(dst) >= BEV_READ_HIGH_WATER) {
+ LOG_DEBUG("stream pipe backpressure pause",
+ "line", _LOGV(conn->route->line_no),
+ "paused", _LOGV("client"),
+ "dst", _LOGV("upstream"),
+ "dst_output_len", _LOGV(evbuffer_get_length(dst))
+ );
+
+ bufferevent_disable(client, EV_READ);
+ }
+}
+
+void pipe_upstream_read_cb(struct bufferevent *upstream, void *arg)
+{
+ conn_t *conn = arg;
+ struct bufferevent *client = conn->client;
+
+ if (client == NULL) {
+ return;
+ }
+
+ struct evbuffer *src = bufferevent_get_input(upstream);
+ struct evbuffer *dst = bufferevent_get_output(client);
+
+#ifdef TINYPROXY_DEBUG
+ size_t len = evbuffer_get_length(src);
+
+ LOG_DEBUG("stream pipe read",
+ "line", _LOGV(conn->route->line_no),
+ "src", _LOGV("upstream"),
+ "dst", _LOGV("client"),
+ "bytes", _LOGV(len)
+ );
+#endif
+
+ evbuffer_add_buffer(dst, src);
+
+ if (evbuffer_get_length(dst) >= BEV_READ_HIGH_WATER) {
+ LOG_DEBUG("stream pipe backpressure pause",
+ "line", _LOGV(conn->route->line_no),
+ "paused", _LOGV("upstream"),
+ "dst", _LOGV("client"),
+ "dst_output_len", _LOGV(evbuffer_get_length(dst))
+ );
+
+ bufferevent_disable(upstream, EV_READ);
+ }
+}
+
+void pipe_client_write_cb(struct bufferevent *client, void *arg)
+{
+ conn_t *conn = arg;
+ struct bufferevent *upstream = conn->upstream;
+
+ struct evbuffer *output = bufferevent_get_output(client);
+ size_t output_len = evbuffer_get_length(output);
+
+ LOG_DEBUG("stream pipe write",
+ "line", _LOGV(conn->route->line_no),
+ "dst", _LOGV("client"),
+ "src", _LOGV("upstream"),
+ "dst_output_len", _LOGV(output_len)
+ );
+
+ if (conn->close_client_after_drain && output_len == 0) {
+ LOG_INFO("client output drained; shutting down write side",
+ "line", _LOGV(conn->route->line_no)
+ );
+
+ finish_client_write(conn);
+ conn->close_client_after_drain = false;
+ conn->close_after_client_eof = true;
+ return;
+ }
+
+ if (upstream != NULL && output_len < BEV_WRITE_RESUME_WATER) {
+ LOG_DEBUG("stream pipe backpressure resume",
+ "line", _LOGV(conn->route->line_no),
+ "resumed", _LOGV("upstream"),
+ "dst", _LOGV("client"),
+ "dst_output_len", _LOGV(output_len)
+ );
+
+ bufferevent_enable(upstream, EV_READ);
+ }
+}
+
+void pipe_upstream_write_cb(struct bufferevent *upstream, void *arg)
+{
+ conn_t *conn = arg;
+ struct bufferevent *client = conn->client;
+
+ struct evbuffer *output = bufferevent_get_output(upstream);
+ size_t output_len = evbuffer_get_length(output);
+
+ LOG_DEBUG("stream pipe write",
+ "line", _LOGV(conn->route->line_no),
+ "dst", _LOGV("upstream"),
+ "src", _LOGV("client"),
+ "dst_output_len", _LOGV(output_len)
+ );
+
+ if (client != NULL && output_len < BEV_WRITE_RESUME_WATER) {
+ LOG_DEBUG("stream pipe backpressure resume",
+ "line", _LOGV(conn->route->line_no),
+ "resumed", _LOGV("client"),
+ "dst", _LOGV("upstream"),
+ "dst_output_len", _LOGV(output_len)
+ );
+
+ bufferevent_enable(client, EV_READ);
+ }
+}
diff --git a/src/stream_pipe.h b/src/stream_pipe.h
new file mode 100644
index 0000000..c8d90fa
--- /dev/null
+++ b/src/stream_pipe.h
@@ -0,0 +1,11 @@
+#ifndef STREAM_PIPE_H
+#define STREAM_PIPE_H
+
+#include <event2/bufferevent.h>
+
+void pipe_client_read_cb(struct bufferevent *client, void *arg);
+void pipe_client_write_cb(struct bufferevent *client, void *arg);
+void pipe_upstream_read_cb(struct bufferevent *upstream, void *arg);
+void pipe_upstream_write_cb(struct bufferevent *upstream, void *arg);
+
+#endif
diff --git a/tests/support.py b/tests/support.py
index c388468..97060c4 100644
--- a/tests/support.py
+++ b/tests/support.py
@@ -308,6 +308,88 @@ def print_process_output(name: str, proc: subprocess.Popen | None) -> None:
print(stderr)
+
+
+@dataclass
+class TrackedStreamServer:
+ server: asyncio.AbstractServer
+ tasks: set[asyncio.Task]
+
+ async def close(self, timeout: float = 3.0) -> None:
+ self.server.close()
+ await self.server.wait_closed()
+
+ if not self.tasks:
+ return
+
+ done, pending = await asyncio.wait(self.tasks, timeout=timeout)
+
+ for task in pending:
+ task.cancel()
+
+ if pending:
+ await asyncio.gather(*pending, return_exceptions=True)
+
+ for task in done:
+ try:
+ task.result()
+ except asyncio.CancelledError:
+ pass
+
+
+async def start_tracked_stream_server(
+ handler,
+ host: str,
+ port: int,
+ backlog: int = 128,
+) -> TrackedStreamServer:
+ tasks: set[asyncio.Task] = set()
+
+ async def tracked_handler(
+ reader: asyncio.StreamReader,
+ writer: asyncio.StreamWriter,
+ ) -> None:
+ task = asyncio.current_task()
+ if task is not None:
+ tasks.add(task)
+
+ try:
+ await handler(reader, writer)
+ finally:
+ if task is not None:
+ tasks.discard(task)
+
+ server = await asyncio.start_server(
+ tracked_handler,
+ host,
+ port,
+ backlog=backlog,
+ )
+
+ return TrackedStreamServer(server=server, tasks=tasks)
+
+
+@asynccontextmanager
+async def run_tracked_stream_server(
+ handler,
+ host: str,
+ port: int,
+ backlog: int = 128,
+ close_timeout: float = 3.0,
+):
+ tracked = await start_tracked_stream_server(
+ handler,
+ host,
+ port,
+ backlog=backlog,
+ )
+
+ try:
+ yield tracked.server
+ finally:
+ await tracked.close(timeout=close_timeout)
+
+
@dataclass
class TinyproxyFixture:
proxy: subprocess.Popen
@@ -321,19 +403,14 @@ async def run_echo_backend(
host: str = LISTEN_HOST,
port: int = BACKEND_PORT,
):
- server = await asyncio.start_server(
+ async with run_tracked_stream_server(
echo_handler,
host,
port,
backlog=4096,
- )
-
- try:
+ ) as server:
await wait_for_port(host, port)
yield server
- finally:
- server.close()
- await server.wait_closed()
@asynccontextmanager
async def run_tinyproxy_with_conf(
diff --git a/tests/test_tcp_backpressure.py b/tests/test_tcp_backpressure.py
index 454eb9f..8b23cbe 100644
--- a/tests/test_tcp_backpressure.py
+++ b/tests/test_tcp_backpressure.py
@@ -7,6 +7,7 @@ from .support import (
PROXY_PORT,
SkipTest,
run_tinyproxy_with_conf,
+ start_tracked_stream_server,
)
BACKPRESSURE_MAX_BYTES = 32 * 1024 * 1024
@@ -93,7 +94,7 @@ async def test_tcp_backpressure_when_upstream_does_not_read() -> None:
finally:
await abort_writer(writer)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
blackhole_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -130,8 +131,7 @@ async def test_tcp_backpressure_when_upstream_does_not_read() -> None:
stop_backend.set()
- backend_server.close()
- await asyncio.wait_for(backend_server.wait_closed(), timeout=3.0)
+ await backend_server.close(timeout=3.0)
async def test_tcp_backpressure_when_client_does_not_read() -> None:
proxy_bin = os.environ.get("TINYPROXY_BIN")
@@ -180,7 +180,7 @@ async def test_tcp_backpressure_when_client_does_not_read() -> None:
finally:
await abort_writer(writer)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
pushing_backend_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -217,8 +217,7 @@ async def test_tcp_backpressure_when_client_does_not_read() -> None:
if writer is not None:
await abort_writer(writer)
- backend_server.close()
- await asyncio.wait_for(backend_server.wait_closed(), timeout=3.0)
+ await backend_server.close(timeout=3.0)
TESTS = [
diff --git a/tests/test_tcp_connect_timeout.py b/tests/test_tcp_connect_timeout.py
index 7f28398..621f4a3 100644
--- a/tests/test_tcp_connect_timeout.py
+++ b/tests/test_tcp_connect_timeout.py
@@ -7,6 +7,7 @@ from .support import (
PROXY_PORT,
SkipTest,
run_tinyproxy_with_conf,
+ start_tracked_stream_server,
)
@@ -102,7 +103,7 @@ async def test_tcp_connect_timeout_does_not_replace_idle_timeout_after_connect()
f" connect_timeout=1,idle_timeout=2\n"
)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -158,8 +159,7 @@ async def test_tcp_connect_timeout_does_not_replace_idle_timeout_after_connect()
finally:
await close_writer(writer)
finally:
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
TESTS = [
diff --git a/tests/test_tcp_idle_timeout.py b/tests/test_tcp_idle_timeout.py
index c3dab6d..ce1a8ab 100644
--- a/tests/test_tcp_idle_timeout.py
+++ b/tests/test_tcp_idle_timeout.py
@@ -7,6 +7,7 @@ from .support import (
PROXY_PORT,
SkipTest,
run_tinyproxy_with_conf,
+ start_tracked_stream_server,
)
@@ -51,7 +52,7 @@ async def test_tcp_idle_timeout_closes_idle_connection() -> None:
f" idle_timeout=1\n"
)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -91,8 +92,7 @@ async def test_tcp_idle_timeout_closes_idle_connection() -> None:
finally:
await close_writer(writer)
finally:
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
async def test_tcp_idle_timeout_keeps_active_connection_open() -> None:
@@ -107,7 +107,7 @@ async def test_tcp_idle_timeout_keeps_active_connection_open() -> None:
f" idle_timeout=2\n"
)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -144,8 +144,7 @@ async def test_tcp_idle_timeout_keeps_active_connection_open() -> None:
finally:
await close_writer(writer)
finally:
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
TESTS = [
diff --git a/tests/test_tcp_keep_alive.py b/tests/test_tcp_keep_alive.py
index 35678e4..9bfffc1 100644
--- a/tests/test_tcp_keep_alive.py
+++ b/tests/test_tcp_keep_alive.py
@@ -12,6 +12,7 @@ from .support import (
PROXY_PORT,
SkipTest,
run_tinyproxy_with_conf,
+ start_tracked_stream_server,
)
@@ -67,7 +68,7 @@ async def test_tcp_keep_alive_roundtrip() -> None:
f" keep_alive,idle_timeout=5\n"
)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -101,8 +102,7 @@ async def test_tcp_keep_alive_roundtrip() -> None:
finally:
await close_writer(writer)
finally:
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
async def test_tcp_keep_alive_sets_socket_option() -> None:
@@ -124,7 +124,7 @@ async def test_tcp_keep_alive_sets_socket_option() -> None:
f" keep_alive,idle_timeout=5\n"
)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -185,8 +185,7 @@ async def test_tcp_keep_alive_sets_socket_option() -> None:
proxy.kill()
proxy.wait(timeout=3.0)
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
trace_text = trace_path.read_text(errors="replace")
diff --git a/tests/test_unix_listeners.py b/tests/test_unix_listeners.py
index c1b3060..7f3624f 100644
--- a/tests/test_unix_listeners.py
+++ b/tests/test_unix_listeners.py
@@ -8,6 +8,8 @@ from .support import (
LISTEN_HOST,
BACKEND_PORT,
SkipTest,
+ echo_handler,
+ start_tracked_stream_server,
run_tinyproxy_with_conf,
)
@@ -41,20 +43,6 @@ def unlink_if_exists(path: str) -> None:
except FileNotFoundError:
pass
-async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
- try:
- while True:
- data = await reader.read(65536)
- if not data:
- break
-
- writer.write(data)
- await writer.drain()
- finally:
- writer.close()
- await writer.wait_closed()
-
-
async def close_writer(writer: asyncio.StreamWriter) -> None:
try:
writer.close()
@@ -187,7 +175,7 @@ async def run_stream_chain_roundtrip(
for path in socks_to_cleanup:
unlink_if_exists(path)
- backend_server = await asyncio.start_server(
+ backend_server = await start_tracked_stream_server(
echo_handler,
LISTEN_HOST,
BACKEND_PORT,
@@ -203,8 +191,7 @@ async def run_stream_chain_roundtrip(
):
return await tcp_roundtrip(LISTEN_HOST, front_port, payload)
finally:
- backend_server.close()
- await backend_server.wait_closed()
+ await backend_server.close()
for path in socks_to_cleanup:
unlink_if_exists(path)