net/proxy/util:write

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",

net/proxy/ares/socket.cc

@@ -2,6 +2,8 @@
 
 #include <utility>
 
+#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<TcpSocket>(this)](
-            std::error_code ec, size_t) {
+        write_buffer_,
+        [socket = boost::intrusive_ptr<TcpSocket>(this)](std::error_code ec) {
             if (ec) {
                 socket->stream_.reset();
             }

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",
     ],
 )

net/proxy/misc/echo-handler.cc

@@ -6,6 +6,7 @@
 #include <utility>
 
 #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> datagram)

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",
     ],

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;

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<TcpConnection>(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<TcpConnection>(this)](
-            std::error_code ec, size_t) {
+            std::error_code ec) {
             if (ec) {
                 connection->close();
                 return;

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",

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<TcpConnection>(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<TcpConnection>(this)](
-            std::error_code ec, size_t) {
+            std::error_code ec) {
             if (ec) {
                 connection->close();
                 return;

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",

net/proxy/system/connector.cc

@@ -4,6 +4,7 @@
 #include <utility>
 
 #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<void(
         std::error_code, std::unique_ptr<Stream>) &&> 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;

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",
+    ],
+)

net/proxy/util/copy.cc

@@ -4,6 +4,7 @@
 #include <utility>
 
 #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;

net/proxy/util/write.cc

@@ -0,0 +1,63 @@
+#include "net/proxy/util/write.h"
+
+#include <utility>
+
+namespace net {
+namespace proxy {
+namespace {
+
+class WriteOperation {
+public:
+    WriteOperation(
+        Stream &stream,
+        ConstBufferSpan buffer,
+        absl::AnyInvocable<void(std::error_code) &&> callback);
+
+    void start() { write(); }
+
+private:
+    void write();
+    void finish(std::error_code ec) { std::move(callback_)(ec); }
+
+    Stream &stream_;
+    ConstBufferSpan buffer_;
+    absl::AnyInvocable<void(std::error_code) &&> callback_;
+};
+
+WriteOperation::WriteOperation(
+    Stream &stream,
+    ConstBufferSpan buffer,
+    absl::AnyInvocable<void(std::error_code) &&> 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<void(std::error_code) &&> callback) {
+    auto *operation = new WriteOperation(stream, buffer, std::move(callback));
+    operation->start();
+}
+
+}  // namespace proxy
+}  // namespace net

net/proxy/util/write.h

@@ -0,0 +1,21 @@
+#ifndef _NET_PROXY_UTIL_WRITE_H
+#define _NET_PROXY_UTIL_WRITE_H
+
+#include <system_error>
+
+#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<void(std::error_code) &&> callback);
+
+}  // namespace proxy
+}  // namespace net
+
+#endif  // _NET_PROXY_UTIL_WRITE_H