net/proxy add datagram support

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

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,13 +75,49 @@ 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) {
     auto *connection = new StreamConnection(std::move(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

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

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,13 +57,35 @@ 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) {
     auto *connection = new StreamConnection(std::move(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

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

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

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

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

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

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)

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;

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

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)

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;

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