From 7da320cc81e110c5e1d150e108aa941aacbb601f Mon Sep 17 00:00:00 2001 From: iceboy Date: Sat, 7 Sep 2024 12:57:01 -0700 Subject: [PATCH] net/proxy/util:write --- net/proxy/ares/BUILD | 1 + net/proxy/ares/socket.cc | 9 +++-- net/proxy/misc/BUILD | 1 + net/proxy/misc/echo-handler.cc | 18 ++++----- net/proxy/shadowsocks/BUILD | 2 + net/proxy/shadowsocks/connector.cc | 8 ++-- net/proxy/shadowsocks/handler.cc | 14 +++---- net/proxy/socks/BUILD | 1 + net/proxy/socks/handler.cc | 13 +++--- net/proxy/system/BUILD | 1 + net/proxy/system/connector.cc | 7 ++-- net/proxy/util/BUILD | 12 ++++++ net/proxy/util/copy.cc | 13 +++--- net/proxy/util/write.cc | 63 ++++++++++++++++++++++++++++++ net/proxy/util/write.h | 21 ++++++++++ 15 files changed, 144 insertions(+), 40 deletions(-) create mode 100644 net/proxy/util/write.cc create mode 100644 net/proxy/util/write.h diff --git a/net/proxy/ares/BUILD b/net/proxy/ares/BUILD index 37bf44e..8636711 100644 --- a/net/proxy/ares/BUILD +++ b/net/proxy/ares/BUILD @@ -32,6 +32,7 @@ cc_library( hdrs = ["socket.h"], deps = [ "//net/proxy:interface", + "//net/proxy/util:write", "//third_party/cares", "@org_boost_boost//:smart_ptr", "@org_iceboy_trunk//net:asio", diff --git a/net/proxy/ares/socket.cc b/net/proxy/ares/socket.cc index 926845e..d7d6286 100644 --- a/net/proxy/ares/socket.cc +++ b/net/proxy/ares/socket.cc @@ -2,6 +2,8 @@ #include +#include "net/proxy/util/write.h" + #ifdef _WIN32 struct iovec { void *iov_base; @@ -162,11 +164,10 @@ ares_ssize_t TcpSocket::sendv(const iovec *data, int len) { SET_ERRNO(ENETUNREACH); return -1; } - async_write( + write( *stream_, - const_buffer(write_buffer_.data(), write_buffer_.size()), - [socket = boost::intrusive_ptr(this)]( - std::error_code ec, size_t) { + write_buffer_, + [socket = boost::intrusive_ptr(this)](std::error_code ec) { if (ec) { socket->stream_.reset(); } diff --git a/net/proxy/misc/BUILD b/net/proxy/misc/BUILD index f4a9f8b..7c7d38d 100644 --- a/net/proxy/misc/BUILD +++ b/net/proxy/misc/BUILD @@ -19,6 +19,7 @@ cc_library( hdrs = ["echo-handler.h"], deps = [ "//net/proxy:interface", + "//net/proxy/util:write", "@com_google_absl//absl/container:fixed_array", ], ) diff --git a/net/proxy/misc/echo-handler.cc b/net/proxy/misc/echo-handler.cc index 526ea2e..0643777 100644 --- a/net/proxy/misc/echo-handler.cc +++ b/net/proxy/misc/echo-handler.cc @@ -6,6 +6,7 @@ #include #include "absl/container/fixed_array.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -63,16 +64,13 @@ void StreamConnection::read() { } void StreamConnection::write() { - async_write( - *stream_, - const_buffer(buffer_.data(), size_), - [this](std::error_code ec, size_t) { - if (ec) { - finish(); - return; - } - read(); - }); + proxy::write(*stream_, {buffer_.data(), size_}, [this](std::error_code ec) { + if (ec) { + finish(); + return; + } + read(); + }); } DatagramConnection::DatagramConnection(std::unique_ptr datagram) diff --git a/net/proxy/shadowsocks/BUILD b/net/proxy/shadowsocks/BUILD index f26cc8d..31c4d6e 100644 --- a/net/proxy/shadowsocks/BUILD +++ b/net/proxy/shadowsocks/BUILD @@ -23,6 +23,7 @@ cc_library( ":pre-shared-key", ":salt-filter", "//net/proxy:interface", + "//net/proxy/util:write", "@com_google_absl//absl/random", "@org_iceboy_trunk//base:logging", "@org_iceboy_trunk//net:endpoint", @@ -65,6 +66,7 @@ cc_library( ":salt-filter", ":session-subkey", "//net/proxy:interface", + "//net/proxy/util:write", "@com_google_absl//absl/container:fixed_array", "@org_iceboy_trunk//base:logging", ], diff --git a/net/proxy/shadowsocks/connector.cc b/net/proxy/shadowsocks/connector.cc index b8595bd..4255d19 100644 --- a/net/proxy/shadowsocks/connector.cc +++ b/net/proxy/shadowsocks/connector.cc @@ -8,6 +8,7 @@ #include "net/proxy/shadowsocks/decryptor.h" #include "net/proxy/shadowsocks/encryptor.h" #include "net/proxy/stream.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -442,12 +443,11 @@ void Connector::TcpStream::write( encryptor_.write_payload_chunk({buffer.data(), buffer.size()}); total_size += buffer.size(); } - ConstBufferSpan write_buffer = encryptor_.buffer(); - async_write( + proxy::write( *base_stream_, - const_buffer(write_buffer.data(), write_buffer.size()), + encryptor_.buffer(), [total_size, callback = std::move(callback)]( - std::error_code ec, size_t) mutable { + std::error_code ec) mutable { if (ec) { std::move(callback)(ec, 0); return; diff --git a/net/proxy/shadowsocks/handler.cc b/net/proxy/shadowsocks/handler.cc index 6051ba2..c001b3a 100644 --- a/net/proxy/shadowsocks/handler.cc +++ b/net/proxy/shadowsocks/handler.cc @@ -11,6 +11,7 @@ #include "base/logging.h" #include "net/proxy/shadowsocks/decryptor.h" #include "net/proxy/shadowsocks/encryptor.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -338,11 +339,11 @@ void Handler::TcpConnection::forward_parse_host(size_t header_length) { } void Handler::TcpConnection::forward_write() { - async_write( + write( *remote_stream_, - const_buffer(decryptor_.pop_buffer(read_length_), read_length_), + {decryptor_.pop_buffer(read_length_), read_length_}, [connection = boost::intrusive_ptr(this)]( - std::error_code ec, size_t) { + std::error_code ec) { if (ec) { connection->close(); return; @@ -391,12 +392,11 @@ void Handler::TcpConnection::backward_write() { encryptor_.write_payload_chunk(read_buffer.subspan(0, chunk_size)); read_buffer.remove_prefix(chunk_size); } while (!read_buffer.empty()); - ConstBufferSpan write_buffer = encryptor_.buffer(); - async_write( + write( *stream_, - const_buffer(write_buffer.data(), write_buffer.size()), + encryptor_.buffer(), [connection = boost::intrusive_ptr(this)]( - std::error_code ec, size_t) { + std::error_code ec) { if (ec) { connection->close(); return; diff --git a/net/proxy/socks/BUILD b/net/proxy/socks/BUILD index 2e9d48d..277729d 100644 --- a/net/proxy/socks/BUILD +++ b/net/proxy/socks/BUILD @@ -18,6 +18,7 @@ cc_library( hdrs = ["handler.h"], deps = [ "//net/proxy:interface", + "//net/proxy/util:write", "@com_google_absl//absl/algorithm", "@com_google_absl//absl/container:fixed_array", "@org_boost_boost//:endian", diff --git a/net/proxy/socks/handler.cc b/net/proxy/socks/handler.cc index e319c7b..b212387 100644 --- a/net/proxy/socks/handler.cc +++ b/net/proxy/socks/handler.cc @@ -9,6 +9,7 @@ #include "absl/container/fixed_array.h" #include "base/logging.h" #include "base/types.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -242,11 +243,11 @@ void Handler::TcpConnection::connect_host(ConstBufferSpan buffer) { } void Handler::TcpConnection::forward_write() { - async_write( + write( *remote_stream_, - const_buffer(forward_buffer_.data(), forward_size_), + {forward_buffer_.data(), forward_size_}, [connection = boost::intrusive_ptr(this)]( - std::error_code ec, size_t) { + std::error_code ec) { if (ec) { connection->close(); return; @@ -272,11 +273,11 @@ void Handler::TcpConnection::reply() { } void Handler::TcpConnection::backward_write() { - async_write( + write( *stream_, - const_buffer(backward_buffer_.data(), backward_size_), + {backward_buffer_.data(), backward_size_}, [connection = boost::intrusive_ptr(this)]( - std::error_code ec, size_t) { + std::error_code ec) { if (ec) { connection->close(); return; diff --git a/net/proxy/system/BUILD b/net/proxy/system/BUILD index 5ea1ca0..98571d8 100644 --- a/net/proxy/system/BUILD +++ b/net/proxy/system/BUILD @@ -21,6 +21,7 @@ cc_library( ":udp-socket-datagram", "//net/proxy:interface", "//net/proxy/ares:resolver", + "//net/proxy/util:write", "@com_google_absl//absl/strings", "@org_iceboy_trunk//net:asio", "@org_iceboy_trunk//net:timer-list", diff --git a/net/proxy/system/connector.cc b/net/proxy/system/connector.cc index 41a4c78..93e0a04 100644 --- a/net/proxy/system/connector.cc +++ b/net/proxy/system/connector.cc @@ -4,6 +4,7 @@ #include #include "absl/strings/str_cat.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -121,11 +122,11 @@ void Connector::send_initial_data( absl::AnyInvocable) &&> callback) { TcpSocketStream &stream_ref = *stream; - async_write( + write( stream_ref, - initial_data, + {initial_data.data(), initial_data.size()}, [stream = std::move(stream), callback = std::move(callback)]( - std::error_code ec, size_t) mutable { + std::error_code ec) mutable { if (ec) { std::move(callback)(ec, nullptr); return; diff --git a/net/proxy/util/BUILD b/net/proxy/util/BUILD index 517cbc2..155dc71 100644 --- a/net/proxy/util/BUILD +++ b/net/proxy/util/BUILD @@ -5,8 +5,20 @@ cc_library( srcs = ["copy.cc"], hdrs = ["copy.h"], deps = [ + ":write", "//net/proxy:interface", "@com_google_absl//absl/container:fixed_array", "@com_google_absl//absl/functional:any_invocable", ], ) + +cc_library( + name = "write", + srcs = ["write.cc"], + hdrs = ["write.h"], + deps = [ + "//net/proxy:interface", + "@com_google_absl//absl/functional:any_invocable", + "@org_iceboy_trunk//base:types", + ], +) diff --git a/net/proxy/util/copy.cc b/net/proxy/util/copy.cc index dc2a7de..167fc3d 100644 --- a/net/proxy/util/copy.cc +++ b/net/proxy/util/copy.cc @@ -4,6 +4,7 @@ #include #include "absl/container/fixed_array.h" +#include "net/proxy/util/write.h" namespace net { namespace proxy { @@ -64,10 +65,10 @@ void CopyBidirOperation::forward_read() { } void CopyBidirOperation::forward_write() { - async_write( + write( *stream1_, - const_buffer(forward_buffer_.data(), forward_size_), - [this](std::error_code ec, size_t) { + {forward_buffer_.data(), forward_size_}, + [this](std::error_code ec) { if (ec) { finish_one(ec); return; @@ -90,10 +91,10 @@ void CopyBidirOperation::backward_read() { } void CopyBidirOperation::backward_write() { - async_write( + write( *stream0_, - const_buffer(backward_buffer_.data(), backward_size_), - [this](std::error_code ec, size_t) { + {backward_buffer_.data(), backward_size_}, + [this](std::error_code ec) { if (ec) { finish_one(ec); return; diff --git a/net/proxy/util/write.cc b/net/proxy/util/write.cc new file mode 100644 index 0000000..b219829 --- /dev/null +++ b/net/proxy/util/write.cc @@ -0,0 +1,63 @@ +#include "net/proxy/util/write.h" + +#include + +namespace net { +namespace proxy { +namespace { + +class WriteOperation { +public: + WriteOperation( + Stream &stream, + ConstBufferSpan buffer, + absl::AnyInvocable callback); + + void start() { write(); } + +private: + void write(); + void finish(std::error_code ec) { std::move(callback_)(ec); } + + Stream &stream_; + ConstBufferSpan buffer_; + absl::AnyInvocable callback_; +}; + +WriteOperation::WriteOperation( + Stream &stream, + ConstBufferSpan buffer, + absl::AnyInvocable callback) + : stream_(stream), + buffer_(buffer), + callback_(std::move(callback)) {} + +void WriteOperation::write() { + stream_.write( + {{buffer_.data(), buffer_.size()}}, + [this](std::error_code ec, size_t size) { + if (ec) { + finish(ec); + return; + } + buffer_.remove_prefix(size); + if (buffer_.empty()) { + finish({}); + return; + } + write(); + }); +} + +} // namespace + +void write( + Stream &stream, + ConstBufferSpan buffer, + absl::AnyInvocable callback) { + auto *operation = new WriteOperation(stream, buffer, std::move(callback)); + operation->start(); +} + +} // namespace proxy +} // namespace net diff --git a/net/proxy/util/write.h b/net/proxy/util/write.h new file mode 100644 index 0000000..5f05bc6 --- /dev/null +++ b/net/proxy/util/write.h @@ -0,0 +1,21 @@ +#ifndef _NET_PROXY_UTIL_WRITE_H +#define _NET_PROXY_UTIL_WRITE_H + +#include + +#include "absl/functional/any_invocable.h" +#include "base/types.h" +#include "net/proxy/stream.h" + +namespace net { +namespace proxy { + +void write( + Stream &stream, + ConstBufferSpan buffer, + absl::AnyInvocable callback); + +} // namespace proxy +} // namespace net + +#endif // _NET_PROXY_UTIL_WRITE_H