Add timeout for tcp connection

Jackarain · Sep 3, 2024 · e983526 · e983526
1 parent 8642b1b
commit e983526
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 1 deletion.
diff --git a/proxy/include/proxy/proxy_server.hpp b/proxy/include/proxy/proxy_server.hpp
@@ -2862,6 +2862,7 @@ R"x*x*x(<html>
 
 			for (; !m_abort;)
 			{
+				stream_expires_after(from, std::chrono::seconds(m_option.tcp_timeout_));
 				auto bytes = co_await from.async_read_some(
 					net::buffer(data), net_awaitable[ec]);
 				if (ec || m_abort)
@@ -2874,6 +2875,7 @@ R"x*x*x(<html>
 					co_return;
 				}
 
+				stream_expires_after(to, std::chrono::seconds(m_option.tcp_timeout_));
 				co_await net::async_write(to,
 					net::buffer(data, bytes), net_awaitable[ec]);
 				if (ec || m_abort)
@@ -4396,6 +4398,75 @@ R"x*x*x(<html>
 			co_return;
 		}
 
+		inline void stream_expires_never(proxy_stream_type& stream)
+		{
+			boost::variant2::visit([](auto& s) mutable
+			{
+				using ValueType = std::decay_t<decltype(s)>;
+				using NextLayerType = util::tcp_socket::next_layer_type;
+
+				if constexpr (std::same_as<NextLayerType, util::tcp_stream>)
+				{
+					if constexpr (std::same_as<util::tcp_socket, ValueType>)
+					{
+						auto& next_layer = s.next_layer();
+						next_layer.expires_never();
+					}
+					else if constexpr (std::same_as<util::ssl_stream, ValueType>)
+					{
+						auto& next_layer = s.next_layer().next_layer();
+						next_layer.expires_never();
+					}
+				}
+			}, stream);
+		}
+
+		inline void stream_expires_after(proxy_stream_type& stream, net::steady_timer::duration expiry_time)
+		{
+			boost::variant2::visit([expiry_time](auto& s) mutable
+			{
+				using ValueType = std::decay_t<decltype(s)>;
+				using NextLayerType = util::tcp_socket::next_layer_type;
+
+				if constexpr (std::same_as<NextLayerType, util::tcp_stream>)
+				{
+					if constexpr (std::same_as<util::tcp_socket, ValueType>)
+					{
+						auto& next_layer = s.next_layer();
+						next_layer.expires_after(expiry_time);
+					}
+					else if constexpr (std::same_as<util::ssl_stream, ValueType>)
+					{
+						auto& next_layer = s.next_layer().next_layer();
+						next_layer.expires_after(expiry_time);
+					}
+				}
+			}, stream);
+		}
+
+		inline void stream_expires_at(proxy_stream_type& stream, net::steady_timer::time_point expiry_time)
+		{
+			boost::variant2::visit([expiry_time](auto& s) mutable
+			{
+				using ValueType = std::decay_t<decltype(s)>;
+				using NextLayerType = util::tcp_socket::next_layer_type;
+
+				if constexpr (std::same_as<NextLayerType, util::tcp_stream>)
+				{
+					if constexpr (std::same_as<util::tcp_socket, ValueType>)
+					{
+						auto& next_layer = s.next_layer();
+						next_layer.expires_at(expiry_time);
+					}
+					else if constexpr (std::same_as<util::ssl_stream, ValueType>)
+					{
+						auto& next_layer = s.next_layer().next_layer();
+						next_layer.expires_at(expiry_time);
+					}
+				}
+			}, stream);
+		}
+
 	private:
 		// m_executor 保存当前 io_context 的 executor.
 		net::any_io_executor m_executor;

diff --git a/server/proxy_server/main.cpp b/server/proxy_server/main.cpp
@@ -95,7 +95,7 @@ bool ssl_prefer_server_ciphers = false;
 int64_t linux_so_mark;
 uint16_t noise_length;
 int udp_timeout;
-
+int tcp_timeout;
 
 //////////////////////////////////////////////////////////////////////////
 
@@ -154,6 +154,7 @@ start_proxy_server(net::io_context& ioc, server_ptr& server)
 	opt.noise_length_ = noise_length;
 	opt.local_ip_ = local_ip;
 	opt.udp_timeout_ = udp_timeout;
+	opt.tcp_timeout_ = tcp_timeout;
 
 	opt.ipip_db_ = ipip_db;
 	if (!deny_region.empty())
@@ -311,6 +312,7 @@ int main(int argc, char** argv)
 		("so_mark", po::value<int64_t>(&linux_so_mark)->default_value(-1), "Set SO_MARK for linux transparent proxy mode.")
 
 		("udp_timeout", po::value<int>(&udp_timeout)->default_value(60), "Set UDP timeout for UDP connections.")
+		("tcp_timeout", po::value<int>(&tcp_timeout)->default_value(60), "Set TCP timeout for TCP connections.")
 
 		("auth_users", po::value<std::vector<std::string>>(&auth_users)->multitoken()->default_value(std::vector<std::string>{"jack:1111"}), "List of authorized users(default user: jack:1111) (e.g: user1:passwd1 user2:passwd2).")