Skip to content

Commit

Permalink
net/proxy support close in stream and datagram
Browse files Browse the repository at this point in the history
  • Loading branch information
iceboy233 committed Mar 31, 2024
1 parent 8ec2174 commit ca875e8
Show file tree
Hide file tree
Showing 11 changed files with 35 additions and 58 deletions.
4 changes: 3 additions & 1 deletion net/proxy/datagram.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ class Datagram {
using executor_type = any_io_executor;

virtual ~Datagram() = default;
virtual any_io_executor get_executor() = 0;

virtual void async_receive_from(
absl::Span<mutable_buffer const> buffers,
Expand All @@ -29,6 +28,9 @@ class Datagram {
const udp::endpoint &endpoint,
absl::AnyInvocable<void(std::error_code, size_t) &&> callback) = 0;

virtual any_io_executor get_executor() = 0;
virtual void close() = 0;

template <typename BuffersT>
void async_receive_from(
const BuffersT &buffers,
Expand Down
8 changes: 1 addition & 7 deletions net/proxy/misc/random-handler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ class StreamConnection : public boost::intrusive_ref_counter<
private:
void read();
void write();
void close() { stream_.reset(); }
void close() { stream_->close(); }

std::unique_ptr<Stream> stream_;
absl::FixedArray<uint8_t, 0> read_buffer_;
Expand Down Expand Up @@ -58,9 +58,6 @@ StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
}

void StreamConnection::read() {
if (!stream_) {
return;
}
stream_->async_read_some(
buffer(read_buffer_.data(), read_buffer_.size()),
[connection = boost::intrusive_ptr<StreamConnection>(this)](
Expand All @@ -74,9 +71,6 @@ void StreamConnection::read() {
}

void StreamConnection::write() {
if (!stream_) {
return;
}
stream_->async_write_some(
const_buffer(write_buffer_.data(), write_buffer_.size()),
[connection = boost::intrusive_ptr<StreamConnection>(this)](
Expand Down
8 changes: 1 addition & 7 deletions net/proxy/misc/zero-handler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ class StreamConnection : public boost::intrusive_ref_counter<
private:
void read();
void write();
void close() { stream_.reset(); }
void close() { stream_->close(); }

std::unique_ptr<Stream> stream_;
absl::FixedArray<uint8_t, 0> read_buffer_;
Expand Down Expand Up @@ -55,9 +55,6 @@ StreamConnection::StreamConnection(std::unique_ptr<Stream> stream)
write_buffer_(8192) {}

void StreamConnection::read() {
if (!stream_) {
return;
}
stream_->async_read_some(
buffer(read_buffer_.data(), read_buffer_.size()),
[connection = boost::intrusive_ptr<StreamConnection>(this)](
Expand All @@ -71,9 +68,6 @@ void StreamConnection::read() {
}

void StreamConnection::write() {
if (!stream_) {
return;
}
stream_->async_write_some(
const_buffer(write_buffer_.data(), write_buffer_.size()),
[connection = boost::intrusive_ptr<StreamConnection>(this)](
Expand Down
7 changes: 2 additions & 5 deletions net/proxy/shadowsocks/connector.cc
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ class Connector::TcpStream : public proxy::Stream {
absl::Span<const_buffer const> buffers,
absl::AnyInvocable<void(std::error_code, size_t) &&> callback) override;

any_io_executor get_executor() override;
any_io_executor get_executor() override { return connector_.executor_; }
void close() override { base_stream_->close(); }

private:
void connect(absl::AnyInvocable<void(std::error_code) &&> callback);
Expand Down Expand Up @@ -455,10 +456,6 @@ void Connector::TcpStream::async_write_some(
});
}

any_io_executor Connector::TcpStream::get_executor() {
return connector_.executor_;
}

} // namespace shadowsocks
} // namespace proxy
} // namespace net
18 changes: 4 additions & 14 deletions net/proxy/shadowsocks/handler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,6 @@ Handler::TcpConnection::TcpConnection(
write_header_(handler_.pre_shared_key_.method().is_spec_2022()) {}

void Handler::TcpConnection::forward_read() {
if (!stream_) {
return;
}
BufferSpan read_buffer = decryptor_.buffer();
stream_->async_read_some(
buffer(read_buffer.data(), read_buffer.size()),
Expand Down Expand Up @@ -342,9 +339,6 @@ void Handler::TcpConnection::forward_parse_host(size_t header_length) {
}

void Handler::TcpConnection::forward_write() {
if (!remote_stream_) {
return;
}
async_write(
*remote_stream_,
buffer(decryptor_.pop_buffer(read_length_), read_length_),
Expand All @@ -361,9 +355,6 @@ void Handler::TcpConnection::forward_write() {
}

void Handler::TcpConnection::backward_read() {
if (!remote_stream_) {
return;
}
remote_stream_->async_read_some(
buffer(backward_read_buffer_.data(), backward_read_buffer_.size()),
[connection = boost::intrusive_ptr<TcpConnection>(this)](
Expand All @@ -378,9 +369,6 @@ void Handler::TcpConnection::backward_read() {
}

void Handler::TcpConnection::backward_write() {
if (!stream_) {
return;
}
ConstBufferSpan read_buffer(
backward_read_buffer_.data(), backward_read_size_);
do {
Expand Down Expand Up @@ -420,8 +408,10 @@ void Handler::TcpConnection::backward_write() {
}

void Handler::TcpConnection::close() {
remote_stream_.reset();
stream_.reset();
if (remote_stream_) {
remote_stream_->close();
}
stream_->close();
}

} // namespace shadowsocks
Expand Down
18 changes: 4 additions & 14 deletions net/proxy/socks/handler.cc
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,6 @@ Handler::TcpConnection::TcpConnection(
backward_buffer_(4096) {}

void Handler::TcpConnection::forward_read() {
if (!stream_) {
return;
}
stream_->async_read_some(
buffer(
&forward_buffer_[forward_size_],
Expand Down Expand Up @@ -246,9 +243,6 @@ void Handler::TcpConnection::connect_host(ConstBufferSpan buffer) {
}

void Handler::TcpConnection::forward_write() {
if (!remote_stream_) {
return;
}
async_write(
*remote_stream_,
buffer(forward_buffer_.data(), forward_size_),
Expand Down Expand Up @@ -279,9 +273,6 @@ void Handler::TcpConnection::reply() {
}

void Handler::TcpConnection::backward_write() {
if (!stream_) {
return;
}
async_write(
*stream_,
buffer(backward_buffer_.data(), backward_size_),
Expand Down Expand Up @@ -309,9 +300,6 @@ void Handler::TcpConnection::backward_dispatch() {
}

void Handler::TcpConnection::backward_read() {
if (!remote_stream_) {
return;
}
remote_stream_->async_read_some(
buffer(
&backward_buffer_[backward_size_],
Expand All @@ -328,8 +316,10 @@ void Handler::TcpConnection::backward_read() {
}

void Handler::TcpConnection::close() {
remote_stream_.reset();
stream_.reset();
if (remote_stream_) {
remote_stream_->close();
}
stream_->close();
}

} // namespace socks
Expand Down
4 changes: 3 additions & 1 deletion net/proxy/stream.h
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ class Stream {
using executor_type = any_io_executor;

virtual ~Stream() = default;
virtual any_io_executor get_executor() = 0;

virtual void async_read_some(
absl::Span<mutable_buffer const> buffers,
Expand All @@ -27,6 +26,9 @@ class Stream {
absl::Span<const_buffer const> buffers,
absl::AnyInvocable<void(std::error_code, size_t) &&> callback) = 0;

virtual any_io_executor get_executor() = 0;
virtual void close() = 0;

template <typename BuffersT>
void async_read_some(
const BuffersT &buffers,
Expand Down
7 changes: 6 additions & 1 deletion net/proxy/system/tcp-socket-stream.cc
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ namespace system {

TcpSocketStream::TcpSocketStream(tcp::socket socket, TimerList &timer_list)
: socket_(std::move(socket)),
timer_(timer_list, [this]() { socket_.close(); }) {}
timer_(timer_list, [this]() { close(); }) {}

void TcpSocketStream::async_read_some(
absl::Span<mutable_buffer const> buffers,
Expand All @@ -28,6 +28,11 @@ void TcpSocketStream::async_write_some(
timer_.update();
}

void TcpSocketStream::close() {
boost::system::error_code ec;
socket_.close(ec);
}

} // namespace proxy
} // namespace system
} // namespace net
7 changes: 3 additions & 4 deletions net/proxy/system/tcp-socket-stream.h
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,6 @@ class TcpSocketStream : public Stream {
TcpSocketStream(const TcpSocketStream &) = delete;
TcpSocketStream &operator=(const TcpSocketStream &) = delete;

any_io_executor get_executor() override {
return socket_.get_executor();
}

void async_read_some(
absl::Span<mutable_buffer const> buffers,
absl::AnyInvocable<void(std::error_code, size_t) &&> callback) override;
Expand All @@ -28,6 +24,9 @@ class TcpSocketStream : public Stream {
absl::Span<const_buffer const> buffers,
absl::AnyInvocable<void(std::error_code, size_t) &&> callback) override;

any_io_executor get_executor() override { return socket_.get_executor(); }
void close() override;

using Stream::async_read_some;
using Stream::async_write_some;

Expand Down
5 changes: 5 additions & 0 deletions net/proxy/system/udp-socket-datagram.cc
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@ void UdpSocketDatagram::async_send_to(
std::move(callback));
}

void UdpSocketDatagram::close() {
boost::system::error_code ec;
socket_.close(ec);
}

} // namespace proxy
} // namespace system
} // namespace net
7 changes: 3 additions & 4 deletions net/proxy/system/udp-socket-datagram.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,6 @@ class UdpSocketDatagram : public Datagram {
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,
Expand All @@ -29,6 +25,9 @@ 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(); }
void close() override;

using Datagram::async_receive_from;
using Datagram::async_send_to;

Expand Down

0 comments on commit ca875e8

Please sign in to comment.