From d48db2e3a32b7784c23c002f48ac787b763e7952 Mon Sep 17 00:00:00 2001
From: iceboy <me@iceboy.org>
Date: Sun, 7 Jan 2024 16:46:36 -0800
Subject: [PATCH] net/proxy add datagram support

---
 net/proxy/handler.h                     |  2 +
 net/proxy/misc/echo-handler.cc          | 53 ++++++++++++++++++++++
 net/proxy/misc/echo-handler.h           |  1 +
 net/proxy/misc/null-handler.cc          | 37 ++++++++++++++++
 net/proxy/misc/null-handler.h           |  1 +
 net/proxy/misc/random-handler.cc        | 58 +++++++++++++++++++++++++
 net/proxy/misc/random-handler.h         |  1 +
 net/proxy/misc/zero-handler.cc          | 55 +++++++++++++++++++++++
 net/proxy/misc/zero-handler.h           |  1 +
 net/proxy/shadowsocks/handler.cc        |  5 +++
 net/proxy/shadowsocks/handler.h         |  1 +
 net/proxy/socks/BUILD                   |  1 +
 net/proxy/socks/handler.cc              |  6 +++
 net/proxy/socks/handler.h               |  1 +
 net/proxy/system/BUILD                  |  1 +
 net/proxy/system/connector.cc           | 12 ++---
 net/proxy/system/listener.cc            | 17 +++++++-
 net/proxy/system/listener.h             |  4 +-
 net/proxy/system/udp-socket-datagram.cc |  3 ++
 net/proxy/system/udp-socket-datagram.h  | 12 ++---
 20 files changed, 258 insertions(+), 14 deletions(-)

diff --git a/net/proxy/handler.h b/net/proxy/handler.h
index 657516d..c1956b8 100644
--- a/net/proxy/handler.h
+++ b/net/proxy/handler.h
@@ -3,6 +3,7 @@
 
 #include <memory>
 
+#include "net/proxy/datagram.h"
 #include "net/proxy/stream.h"
 
 namespace net {
@@ -13,6 +14,7 @@ class Handler {
     virtual ~Handler() = default;
 
     virtual void handle_stream(std::unique_ptr<Stream> stream) = 0;
+    virtual void handle_datagram(std::unique_ptr<Datagram> datagram) = 0;
 };
 
 }  // namespace proxy
diff --git a/net/proxy/misc/echo-handler.cc b/net/proxy/misc/echo-handler.cc
index fef5fa2..5ee4981 100644
--- a/net/proxy/misc/echo-handler.cc
+++ b/net/proxy/misc/echo-handler.cc
@@ -28,6 +28,23 @@ class StreamConnection {
     size_t size_;
 };
 
+class DatagramConnection {
+public:
+    explicit DatagramConnection(std::unique_ptr<Datagram> datagram);
+
+    void start() { read(); }
+
+private:
+    void read();
+    void write();
+    void finish() { delete this; }
+
+    std::unique_ptr<Datagram> datagram_;
+    absl::FixedArray<uint8_t, 0> buffer_;
+    udp::endpoint endpoint_;
+    size_t size_;
+};
+
 StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
     : stream_(std::move(stream)),
       buffer_(8192) {}
@@ -58,6 +75,37 @@ void StreamConnection::write() {
         });
 }
 
+DatagramConnection::DatagramConnection(std::unique_ptr<Datagram> datagram)
+    : datagram_(std::move(datagram)),
+      buffer_(8192) {}
+
+void DatagramConnection::read() {
+    datagram_->async_receive_from(
+        buffer(buffer_.data(), buffer_.size()),
+        endpoint_,
+        [this](std::error_code ec, size_t size) {
+            if (ec) {
+                finish();
+                return;
+            }
+            size_ = size;
+            write();
+        });
+}
+
+void DatagramConnection::write() {
+    datagram_->async_send_to(
+        const_buffer(buffer_.data(), size_),
+        endpoint_,
+        [this](std::error_code ec, size_t) {
+            if (ec) {
+                finish();
+                return;
+            }
+            read();
+        });
+}
+
 }  // namespace
 
 void EchoHandler::handle_stream(std::unique_ptr<Stream> stream) {
@@ -65,6 +113,11 @@ void EchoHandler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void EchoHandler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    auto *connection = new DatagramConnection(std::move(datagram));
+    connection->start();
+}
+
 }  // namespace misc
 }  // namespace proxy
 }  // namespace net
diff --git a/net/proxy/misc/echo-handler.h b/net/proxy/misc/echo-handler.h
index a55d1c7..b2db891 100644
--- a/net/proxy/misc/echo-handler.h
+++ b/net/proxy/misc/echo-handler.h
@@ -10,6 +10,7 @@ namespace misc {
 class EchoHandler : public Handler {
 public:
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 };
 
 }  // namespace misc
diff --git a/net/proxy/misc/null-handler.cc b/net/proxy/misc/null-handler.cc
index f7337fb..c400933 100644
--- a/net/proxy/misc/null-handler.cc
+++ b/net/proxy/misc/null-handler.cc
@@ -26,6 +26,21 @@ class StreamConnection {
     absl::FixedArray<uint8_t, 0> buffer_;
 };
 
+class DatagramConnection {
+public:
+    explicit DatagramConnection(std::unique_ptr<Datagram> datagram);
+
+    void start() { read(); }
+
+private:
+    void read();
+    void finish() { delete this; }
+
+    std::unique_ptr<Datagram> datagram_;
+    absl::FixedArray<uint8_t, 0> buffer_;
+    udp::endpoint endpoint_;
+};
+
 StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
     : stream_(std::move(stream)),
       buffer_(8192) {}
@@ -42,6 +57,23 @@ void StreamConnection::read() {
         });
 }
 
+DatagramConnection::DatagramConnection(std::unique_ptr<Datagram> datagram)
+    : datagram_(std::move(datagram)),
+      buffer_(8192) {}
+
+void DatagramConnection::read() {
+    datagram_->async_receive_from(
+        buffer(buffer_.data(), buffer_.size()),
+        endpoint_,
+        [this](std::error_code ec, size_t) {
+            if (ec) {
+                finish();
+                return;
+            }
+            read();
+        });
+}
+
 }  // namespace
 
 void NullHandler::handle_stream(std::unique_ptr<Stream> stream) {
@@ -49,6 +81,11 @@ void NullHandler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void NullHandler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    auto *connection = new DatagramConnection(std::move(datagram));
+    connection->start();
+}
+
 }  // namespace misc
 }  // namespace proxy
 }  // namespace net
diff --git a/net/proxy/misc/null-handler.h b/net/proxy/misc/null-handler.h
index 6c76ef2..2d4bfac 100644
--- a/net/proxy/misc/null-handler.h
+++ b/net/proxy/misc/null-handler.h
@@ -10,6 +10,7 @@ namespace misc {
 class NullHandler : public Handler {
 public:
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 };
 
 }  // namespace misc
diff --git a/net/proxy/misc/random-handler.cc b/net/proxy/misc/random-handler.cc
index f368da9..87eb29d 100644
--- a/net/proxy/misc/random-handler.cc
+++ b/net/proxy/misc/random-handler.cc
@@ -32,6 +32,24 @@ class StreamConnection : public boost::intrusive_ref_counter<
     absl::FixedArray<uint8_t, 0> write_buffer_;
 };
 
+class DatagramConnection {
+public:
+    explicit DatagramConnection(std::unique_ptr<Datagram> datagram);
+
+    void start() { read(); }
+
+private:
+    void read();
+    void write();
+    void finish() { delete this; }
+
+    std::unique_ptr<Datagram> datagram_;
+    absl::FixedArray<uint8_t, 0> read_buffer_;
+    absl::FixedArray<uint8_t, 0> write_buffer_;
+    udp::endpoint endpoint_;
+    size_t size_;
+};
+
 StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
     : stream_(std::move(stream)),
       read_buffer_(8192),
@@ -72,6 +90,41 @@ void StreamConnection::write() {
         });
 }
 
+DatagramConnection::DatagramConnection(std::unique_ptr<Datagram> datagram)
+    : datagram_(std::move(datagram)),
+      read_buffer_(8192),
+      write_buffer_(8192) {
+    RAND_bytes(write_buffer_.data(), write_buffer_.size());
+}
+
+void DatagramConnection::read() {
+    datagram_->async_receive_from(
+        buffer(read_buffer_.data(), read_buffer_.size()),
+        endpoint_,
+        [this](std::error_code ec, size_t size) {
+            if (ec) {
+                finish();
+                return;
+            }
+            size_ = size;
+            write();
+        });
+}
+
+void DatagramConnection::write() {
+    datagram_->async_send_to(
+        const_buffer(write_buffer_.data(), size_),
+        endpoint_,
+        [this](std::error_code ec, size_t size) {
+            if (ec) {
+                finish();
+                return;
+            }
+            RAND_bytes(write_buffer_.data(), size);
+            read();
+        });
+}
+
 }  // namespace
 
 void RandomHandler::handle_stream(std::unique_ptr<Stream> stream) {
@@ -80,6 +133,11 @@ void RandomHandler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void RandomHandler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    auto *connection = new DatagramConnection(std::move(datagram));
+    connection->start();
+}
+
 }  // namespace misc
 }  // namespace proxy
 }  // namespace net
diff --git a/net/proxy/misc/random-handler.h b/net/proxy/misc/random-handler.h
index 4fa1cfd..b055d9f 100644
--- a/net/proxy/misc/random-handler.h
+++ b/net/proxy/misc/random-handler.h
@@ -10,6 +10,7 @@ namespace misc {
 class RandomHandler : public Handler {
 public:
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 };
 
 }  // namespace misc
diff --git a/net/proxy/misc/zero-handler.cc b/net/proxy/misc/zero-handler.cc
index 00eec12..10b30a7 100644
--- a/net/proxy/misc/zero-handler.cc
+++ b/net/proxy/misc/zero-handler.cc
@@ -31,6 +31,24 @@ class StreamConnection : public boost::intrusive_ref_counter<
     absl::FixedArray<uint8_t, 0> write_buffer_;
 };
 
+class DatagramConnection {
+public:
+    explicit DatagramConnection(std::unique_ptr<Datagram> datagram);
+
+    void start() { read(); }
+
+private:
+    void read();
+    void write();
+    void finish() { delete this; }
+
+    std::unique_ptr<Datagram> datagram_;
+    absl::FixedArray<uint8_t, 0> read_buffer_;
+    absl::FixedArray<uint8_t, 0> write_buffer_;
+    udp::endpoint endpoint_;
+    size_t size_;
+};
+
 StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
     : stream_(std::move(stream)),
       read_buffer_(8192),
@@ -68,6 +86,38 @@ void StreamConnection::write() {
         });
 }
 
+DatagramConnection::DatagramConnection(std::unique_ptr<Datagram> datagram)
+    : datagram_(std::move(datagram)),
+      read_buffer_(8192),
+      write_buffer_(8192) {}
+
+void DatagramConnection::read() {
+    datagram_->async_receive_from(
+        buffer(read_buffer_.data(), read_buffer_.size()),
+        endpoint_,
+        [this](std::error_code ec, size_t size) {
+            if (ec) {
+                finish();
+                return;
+            }
+            size_ = size;
+            write();
+        });
+}
+
+void DatagramConnection::write() {
+    datagram_->async_send_to(
+        const_buffer(write_buffer_.data(), size_),
+        endpoint_,
+        [this](std::error_code ec, size_t) {
+            if (ec) {
+                finish();
+                return;
+            }
+            read();
+        });
+}
+
 }  // namespace
 
 void ZeroHandler::handle_stream(std::unique_ptr<Stream> stream) {
@@ -76,6 +126,11 @@ void ZeroHandler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void ZeroHandler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    auto *connection = new DatagramConnection(std::move(datagram));
+    connection->start();
+}
+
 }  // namespace misc
 }  // namespace proxy
 }  // namespace net
diff --git a/net/proxy/misc/zero-handler.h b/net/proxy/misc/zero-handler.h
index ecffcd5..ea90572 100644
--- a/net/proxy/misc/zero-handler.h
+++ b/net/proxy/misc/zero-handler.h
@@ -10,6 +10,7 @@ namespace misc {
 class ZeroHandler : public Handler {
 public:
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 };
 
 }  // namespace misc
diff --git a/net/proxy/shadowsocks/handler.cc b/net/proxy/shadowsocks/handler.cc
index 03264cc..fefd8ad 100644
--- a/net/proxy/shadowsocks/handler.cc
+++ b/net/proxy/shadowsocks/handler.cc
@@ -71,6 +71,11 @@ void Handler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void Handler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    // TODO: support datagram
+    LOG(warning) << "datagram is not supported yet";
+}
+
 Handler::TcpConnection::TcpConnection(
     Handler &handler,
     std::unique_ptr<Stream> stream)
diff --git a/net/proxy/shadowsocks/handler.h b/net/proxy/shadowsocks/handler.h
index 38af743..95496bd 100644
--- a/net/proxy/shadowsocks/handler.h
+++ b/net/proxy/shadowsocks/handler.h
@@ -25,6 +25,7 @@ class Handler : public proxy::Handler {
     bool init(const InitOptions &config);
 
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 
 private:
     class TcpConnection;
diff --git a/net/proxy/socks/BUILD b/net/proxy/socks/BUILD
index e625279..2e9d48d 100644
--- a/net/proxy/socks/BUILD
+++ b/net/proxy/socks/BUILD
@@ -21,6 +21,7 @@ cc_library(
         "@com_google_absl//absl/algorithm",
         "@com_google_absl//absl/container:fixed_array",
         "@org_boost_boost//:endian",
+        "@org_iceboy_trunk//base:logging",
         "@org_iceboy_trunk//base:types",
         "@org_iceboy_trunk//net:asio",
     ],
diff --git a/net/proxy/socks/handler.cc b/net/proxy/socks/handler.cc
index 0fe33e1..8e472f2 100644
--- a/net/proxy/socks/handler.cc
+++ b/net/proxy/socks/handler.cc
@@ -7,6 +7,7 @@
 
 #include "absl/algorithm/algorithm.h"
 #include "absl/container/fixed_array.h"
+#include "base/logging.h"
 #include "base/types.h"
 
 namespace net {
@@ -60,6 +61,11 @@ void Handler::handle_stream(std::unique_ptr<Stream> stream) {
     connection->start();
 }
 
+void Handler::handle_datagram(std::unique_ptr<Datagram> datagram) {
+    // TODO: Report support status to avoid this being called.
+    LOG(warning) << "datagram is not supported";
+}
+
 Handler::TcpConnection::TcpConnection(
     Handler &handler,
     std::unique_ptr<Stream> stream)
diff --git a/net/proxy/socks/handler.h b/net/proxy/socks/handler.h
index cfc224c..c10f21c 100644
--- a/net/proxy/socks/handler.h
+++ b/net/proxy/socks/handler.h
@@ -14,6 +14,7 @@ class Handler : public proxy::Handler {
     Handler(const any_io_executor &executor, proxy::Connector &connector);
 
     void handle_stream(std::unique_ptr<Stream> stream) override;
+    void handle_datagram(std::unique_ptr<Datagram> datagram) override;
 
 private:
     class TcpConnection;
diff --git a/net/proxy/system/BUILD b/net/proxy/system/BUILD
index 6136b0e..2da7213 100644
--- a/net/proxy/system/BUILD
+++ b/net/proxy/system/BUILD
@@ -33,6 +33,7 @@ cc_library(
     hdrs = ["listener.h"],
     deps = [
         ":tcp-socket-stream",
+        ":udp-socket-datagram",
         "//net/proxy:interface",
         "@org_iceboy_trunk//base:logging",
         "@org_iceboy_trunk//net:asio",
diff --git a/net/proxy/system/connector.cc b/net/proxy/system/connector.cc
index 50ef753..41a4c78 100644
--- a/net/proxy/system/connector.cc
+++ b/net/proxy/system/connector.cc
@@ -63,24 +63,24 @@ void Connector::connect_tcp_host(
 }
 
 std::error_code Connector::bind_udp_v4(std::unique_ptr<Datagram> &datagram) {
-    auto new_datagram = std::make_unique<UdpSocketDatagram>(executor_);
+    udp::socket socket(executor_);
     boost::system::error_code ec;
-    new_datagram->socket().open(udp::v4(), ec);
+    socket.open(udp::v4(), ec);
     if (ec) {
         return ec;
     }
-    datagram = std::move(new_datagram);
+    datagram = std::make_unique<UdpSocketDatagram>(std::move(socket));
     return {};
 }
 
 std::error_code Connector::bind_udp_v6(std::unique_ptr<Datagram> &datagram) {
-    auto new_datagram = std::make_unique<UdpSocketDatagram>(executor_);
+    udp::socket socket(executor_);
     boost::system::error_code ec;
-    new_datagram->socket().open(udp::v6(), ec);
+    socket.open(udp::v6(), ec);
     if (ec) {
         return ec;
     }
-    datagram = std::move(new_datagram);
+    datagram = std::make_unique<UdpSocketDatagram>(std::move(socket));
     return {};
 }
 
diff --git a/net/proxy/system/listener.cc b/net/proxy/system/listener.cc
index 8d956ef..c13255f 100644
--- a/net/proxy/system/listener.cc
+++ b/net/proxy/system/listener.cc
@@ -4,6 +4,7 @@
 
 #include "base/logging.h"
 #include "net/proxy/system/tcp-socket-stream.h"
+#include "net/proxy/system/udp-socket-datagram.h"
 
 namespace net {
 namespace proxy {
@@ -15,12 +16,18 @@ Listener::Listener(
     Handler &handler,
     const Options &options)
     : executor_(executor),
-      tcp_acceptor_(executor, endpoint),
+      endpoint_(endpoint),
       handler_(handler),
+      tcp_acceptor_(executor_, endpoint),
       timer_list_(executor_, options.timeout),
       tcp_no_delay_(options.tcp_no_delay),
       accept_error_delay_(options.accept_error_delay),
-      accept_error_timer_(executor_) { accept(); }
+      accept_error_timer_(executor_) {
+    // TODO: Support more flexible config, such as enabling TCP or UDP
+    // individually, using different handlers, and live config reloading.
+    accept();
+    bind();
+}
 
 void Listener::accept() {
     tcp_acceptor_.async_accept(
@@ -51,6 +58,12 @@ void Listener::accept_error_wait() {
     });
 }
 
+void Listener::bind() {
+    udp::socket socket(executor_, endpoint_);
+    handler_.handle_datagram(
+        std::make_unique<UdpSocketDatagram>(std::move(socket)));
+}
+
 }  // namespace system
 }  // namespace proxy
 }  // namespace net
diff --git a/net/proxy/system/listener.h b/net/proxy/system/listener.h
index 5cf28d9..b7883a0 100644
--- a/net/proxy/system/listener.h
+++ b/net/proxy/system/listener.h
@@ -30,10 +30,12 @@ class Listener {
 private:
     void accept();
     void accept_error_wait();
+    void bind();
 
     any_io_executor executor_;
-    tcp::acceptor tcp_acceptor_;
+    Endpoint endpoint_;
     Handler &handler_;
+    tcp::acceptor tcp_acceptor_;
     TimerList timer_list_;
     bool tcp_no_delay_;
     std::chrono::nanoseconds accept_error_delay_;
diff --git a/net/proxy/system/udp-socket-datagram.cc b/net/proxy/system/udp-socket-datagram.cc
index 48d2a7c..3859ccf 100644
--- a/net/proxy/system/udp-socket-datagram.cc
+++ b/net/proxy/system/udp-socket-datagram.cc
@@ -6,6 +6,9 @@ namespace net {
 namespace proxy {
 namespace system {
 
+UdpSocketDatagram::UdpSocketDatagram(udp::socket socket)
+    : socket_(std::move(socket)) {}
+
 void UdpSocketDatagram::async_receive_from(
     absl::Span<mutable_buffer const> buffers,
     udp::endpoint &endpoint,
diff --git a/net/proxy/system/udp-socket-datagram.h b/net/proxy/system/udp-socket-datagram.h
index 0b60314..6a81b5a 100644
--- a/net/proxy/system/udp-socket-datagram.h
+++ b/net/proxy/system/udp-socket-datagram.h
@@ -10,12 +10,15 @@ namespace system {
 
 class UdpSocketDatagram : public Datagram {
 public:
-    explicit UdpSocketDatagram(const any_io_executor &executor)
-        : socket_(executor) {}
+    explicit UdpSocketDatagram(udp::socket socket);
 
     UdpSocketDatagram(const UdpSocketDatagram &) = delete;
     UdpSocketDatagram &operator=(const UdpSocketDatagram &) = delete;
 
+    any_io_executor get_executor() override {
+        return socket_.get_executor();
+    }
+
     void async_receive_from(
         absl::Span<mutable_buffer const> buffers,
         udp::endpoint &endpoint,
@@ -26,9 +29,8 @@ class UdpSocketDatagram : public Datagram {
         const udp::endpoint &endpoint,
         absl::AnyInvocable<void(std::error_code, size_t) &&> callback) override;
 
-    any_io_executor get_executor() override {
-        return socket_.get_executor();
-    }
+    using Datagram::async_receive_from;
+    using Datagram::async_send_to;
 
     udp::socket &socket() { return socket_; }
     const udp::socket &socket() const { return socket_; }