From bcd4478f44d97169bb79f88983ee2e04b540b753 Mon Sep 17 00:00:00 2001
From: jackarain <jack.wgm@gmail.com>
Date: Wed, 18 Oct 2023 14:24:21 +0800
Subject: [PATCH] Add disable socks proxy protocol.

---
 proxy/include/proxy/proxy_server.hpp | 45 +++++++++++++++++++++++++++-
 server/proxy_server/main.cpp         |  4 +++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/proxy/include/proxy/proxy_server.hpp b/proxy/include/proxy/proxy_server.hpp
index 90a921e995..21e138bc21 100644
--- a/proxy/include/proxy/proxy_server.hpp
+++ b/proxy/include/proxy/proxy_server.hpp
@@ -179,8 +179,11 @@ namespace proxy {
 		// http doc 目录, 用于伪装成web站点.
 		std::string doc_directory_;
 
-		// 禁用未加密 http 服务.
+		// 禁用 http 服务.
 		bool disable_http_{ false };
+
+		// 禁用 socks proxy 服务.
+		bool disable_socks_{ false };
 	};
 
 	// proxy server 虚基类, 任何 proxy server 的实现, 必须基于这个基类.
@@ -333,6 +336,14 @@ namespace proxy {
 
 			if (socks_version == SOCKS_VERSION_5)
 			{
+				if (m_option.disable_socks_)
+				{
+					LOG_DBG << "socks protocol"
+						", connection id: " << m_connection_id
+						<< ", Forbidden";
+					co_return;
+				}
+
 				LOG_DBG << "connection id: " << m_connection_id
 					<< ", socks version: " << socks_version;
 
@@ -341,6 +352,14 @@ namespace proxy {
 			}
 			if (socks_version == SOCKS_VERSION_4)
 			{
+				if (m_option.disable_socks_)
+				{
+					LOG_DBG << "socks protocol"
+						", connection id: " << m_connection_id
+						<< ", Forbidden";
+					co_return;
+				}
+
 				LOG_DBG << "connection id: " << m_connection_id
 					<< ", socks version: " << socks_version;
 
@@ -349,6 +368,14 @@ namespace proxy {
 			}
 			if (socks_version == 'G')
 			{
+				if (m_option.disable_http_)
+				{
+					LOG_DBG << "http protocol"
+						", connection id: " << m_connection_id
+						<< ", Forbidden";
+					co_return;
+				}
+
 				auto ret = co_await http_proxy_get();
 				if (!ret)
 				{
@@ -361,6 +388,14 @@ namespace proxy {
 			}
 			else if (socks_version == 'C')
 			{
+				if (m_option.disable_http_)
+				{
+					LOG_DBG << "http protocol"
+						", connection id: " << m_connection_id
+						<< ", Forbidden";
+					co_return;
+				}
+
 				auto ret = co_await http_proxy_connect();
 				if (!ret)
 				{
@@ -2911,6 +2946,14 @@ Content-Length: 0
 				// plain socks4/5 protocol.
 				if (detect[0] == 0x05 || detect[0] == 0x04)
 				{
+					if (m_option.disable_socks_)
+					{
+						LOG_DBG << "socks protocol"
+							", connection id: " << connection_id
+							<< ", Forbidden";
+						continue;
+					}
+
 					LOG_DBG << "socks protocol:"
 						" connection id: " << connection_id;
 
diff --git a/server/proxy_server/main.cpp b/server/proxy_server/main.cpp
index ab705d29b1..7a5579caed 100644
--- a/server/proxy_server/main.cpp
+++ b/server/proxy_server/main.cpp
@@ -133,6 +133,7 @@ std::string socks_listen;
 std::string doc_directory;
 std::string log_directory;
 bool disable_http = false;
+bool disable_socks = false;
 bool disable_logs;
 bool reuse_port = false;
 
@@ -184,6 +185,8 @@ start_proxy_server(net::io_context& ioc, server_ptr& server)
 	opt.ssl_sni_ = ssl_sni;
 
 	opt.disable_http_ = disable_http;
+	opt.disable_socks_ = disable_socks;
+
 	opt.reuse_port_ = reuse_port;
 
 	opt.doc_directory_ = doc_directory;
@@ -332,6 +335,7 @@ int main(int argc, char** argv)
 		("logs_path", po::value<std::string>(&log_directory)->value_name(""), "Logs dirctory.")
 		("disable_logs", po::value<bool>(&disable_logs)->value_name(""), "Disable logs.")
 		("disable_http", po::value<bool>(&disable_http)->value_name("")->default_value(false), "Disable http protocol.")
+		("disable_socks", po::value<bool>(&disable_socks)->value_name("")->default_value(false), "Disable socks proxy protocol.")
 	;
 
 	// 解析命令行.