From 6ab8af1ef531601f6072ae03eff755ff5d6a49f4 Mon Sep 17 00:00:00 2001 From: myl7 Date: Tue, 28 Dec 2021 16:33:35 +0800 Subject: [PATCH] Refactor SOCKS socket code Separate sending cmd and receiving reply out as they are shared in procedures. Change returns of start* method. --- src/DHTConnectionSocksProxyImpl.cc | 14 ++--- src/SocksProxySocket.cc | 83 ++++++++++++++++++------------ src/SocksProxySocket.h | 34 ++++++------ 3 files changed, 75 insertions(+), 56 deletions(-) diff --git a/src/DHTConnectionSocksProxyImpl.cc b/src/DHTConnectionSocksProxyImpl.cc index 16f69b9c..45202c19 100644 --- a/src/DHTConnectionSocksProxyImpl.cc +++ b/src/DHTConnectionSocksProxyImpl.cc @@ -63,15 +63,15 @@ bool DHTConnectionSocksProxyImpl::startProxy(const std::string& host, // Authentication negotiation bool noAuth = user.empty() || passwd.empty(); if (noAuth) { - int authMethod = - socket_->negotiateAuth(std::vector{SOCKS_AUTH_NO_AUTH}); + int authMethod = socket_->negotiateAuth( + std::vector{SOCKS_AUTH_NO_AUTH}); if (authMethod < 0) { return false; } } else { - int authMethod = socket_->negotiateAuth( - std::vector{SOCKS_AUTH_NO_AUTH, SOCKS_AUTH_USERPASS}); + int authMethod = socket_->negotiateAuth(std::vector{ + SOCKS_AUTH_NO_AUTH, SOCKS_AUTH_USERPASS}); if (authMethod < 0) { return false; } @@ -86,9 +86,9 @@ bool DHTConnectionSocksProxyImpl::startProxy(const std::string& host, } // UDP associate - ssize_t i = socket_->startUdpAssociate(listenAddr, listenPort, - std::make_pair(&bndAddr_, &bndPort_)); - if (i < 0) { + int i = + socket_->startUdpAssociate(listenAddr, listenPort, bndAddr_, bndPort_); + if (i != 0) { return false; } return true; diff --git a/src/SocksProxySocket.cc b/src/SocksProxySocket.cc index 075d9ff1..ebb18b7e 100644 --- a/src/SocksProxySocket.cc +++ b/src/SocksProxySocket.cc @@ -45,15 +45,17 @@ const char C_SOCKS_VER = '\x05'; const char C_AUTH_NONE = '\xff'; const char C_AUTH_USERPASS_VER = '\x01'; const char C_AUTH_USERPASS_OK = '\x00'; -const char C_CMD_UDP_ASSOCIATE = '\x03'; -const char C_ADDR_INET = '\x01'; -const char C_ADDR_DOMAIN = '\x03'; -const char C_ADDR_INET6 = '\x04'; const char C_OK = '\x00'; } // namespace namespace aria2 { +enum SocksProxyAddrType { + SOCKS_ADDR_INET = 1, + SOCKS_ADDR_DOMAIN = 3, + SOCKS_ADDR_INET6 = 4, +}; + SocksProxySocket::SocksProxySocket(int family) : family_(family) {} SocksProxySocket::~SocksProxySocket() = default; @@ -70,13 +72,13 @@ void SocksProxySocket::establish(std::unique_ptr socket) socket_ = std::move(socket); } -int SocksProxySocket::negotiateAuth(std::vector expected) +int SocksProxySocket::negotiateAuth(std::vector expected) { std::stringstream req; req << C_SOCKS_VER; req << static_cast(expected.size()); for (auto c : expected) { - req << c; + req << static_cast(c); } socket_->writeData(req.str()); @@ -109,77 +111,81 @@ char SocksProxySocket::authByUserpass(const std::string& user, return res[1]; } -ssize_t -SocksProxySocket::startUdpAssociate(const std::string& listenAddr, - uint16_t listenPort, - std::pair bnd) +void SocksProxySocket::sendCmd(SocksProxyCmd cmd, const std::string& dstAddr, + uint16_t dstPort, bool allowEmpty) { std::stringstream req; - req << C_SOCKS_VER << C_CMD_UDP_ASSOCIATE << '\x00'; + req << C_SOCKS_VER << static_cast(cmd) << '\x00'; if (family_ == AF_INET) { - if (!listenAddr.empty()) { + if (!allowEmpty || !dstAddr.empty()) { char addrBuf[10]; - net::getBinAddr(addrBuf, listenAddr); - req << C_ADDR_INET << std::string(addrBuf, 4); + net::getBinAddr(addrBuf, dstAddr); + req << static_cast(SOCKS_ADDR_INET) << std::string(addrBuf, 4); } else { req << std::string(4, '\x00'); } } else { - if (!listenAddr.empty()) { + if (!allowEmpty || !dstAddr.empty()) { char addrBuf[20]; - net::getBinAddr(addrBuf, listenAddr); - req << C_ADDR_INET6 << std::string(addrBuf, 16); + net::getBinAddr(addrBuf, dstAddr); + req << static_cast(SOCKS_ADDR_INET6) << std::string(addrBuf, 16); } else { req << std::string(16, '\x00'); } } - if (listenPort) { - uint16_t listenPortBuf = htons(listenPort); + if (dstPort) { + uint16_t listenPortBuf = htons(dstPort); req << std::string(reinterpret_cast(&listenPortBuf), 2); } else { req << std::string(2, '\x00'); } socket_->writeData(req.str()); +} +int SocksProxySocket::receiveReply(int& bndFamily, std::string& bndAddr, + uint16_t& bndPort) +{ char res[5]; size_t resLen = sizeof(res); socket_->readData(res, resLen); - if (res[1] != C_OK) { + int rep = res[1]; + if (rep != C_OK) { socket_->closeConnection(); - return -1; + return rep; } - std::string bndAddr; - uint16_t bndPort; - if (res[3] == C_ADDR_INET) { + if (res[3] == SOCKS_ADDR_INET) { char addrBuf[6]; addrBuf[0] = res[4]; size_t addrLen = sizeof(addrBuf) - 1; socket_->readData(addrBuf + 1, addrLen); char addrStrBuf[20]; inetNtop(AF_INET, addrBuf, addrStrBuf, 20); + bndFamily = AF_INET; bndAddr = std::string(addrStrBuf); bndPort = ntohs(*reinterpret_cast(addrBuf + 4)); } - else if (res[3] == C_ADDR_INET6) { + else if (res[3] == SOCKS_ADDR_INET6) { char addrBuf[18]; addrBuf[0] = res[4]; size_t addrLen = sizeof(addrBuf) - 1; socket_->readData(addrBuf + 1, addrLen); char addrStrBuf[50]; inetNtop(AF_INET6, addrBuf, addrStrBuf, 50); + bndFamily = AF_INET6; bndAddr = std::string(addrStrBuf); bndPort = ntohs(*reinterpret_cast(addrBuf + 16)); } - else if (res[3] == C_ADDR_DOMAIN) { + else if (res[3] == SOCKS_ADDR_DOMAIN) { // 2 more bytes to hold port temporarily. size_t resLen = res[4] + 2; bndAddr = std::string(resLen, '\x00'); socket_->readData(&bndAddr[0], resLen); + bndFamily = AF_INET + AF_INET6; bndPort = ntohs(*reinterpret_cast(&bndAddr[0] + res[4])); bndAddr.resize(res[4]); } @@ -187,15 +193,26 @@ SocksProxySocket::startUdpAssociate(const std::string& listenAddr, socket_->closeConnection(); return -1; } + return rep; +} - ssize_t i = static_cast(bndAddrs_.size()); - bndAddrs_.push_back(bndAddr); - bndPorts_.push_back(bndPort); - if (bnd.first && bnd.second) { - *(bnd.first) = bndAddr; - *(bnd.second) = bndPort; +int SocksProxySocket::startUdpAssociate(const std::string& listenAddr, + uint16_t listenPort, + std::string& bndAddr, uint16_t& bndPort) +{ + sendCmd(SOCKS_CMD_UDP_ASSOCIATE, listenAddr, listenPort, true); + + int bFamily; + std::string bAddr; + uint16_t bPort; + int rep = receiveReply(bFamily, bAddr, bPort); + if (rep != C_OK) { + return rep; } - return i; + + bndAddr = bAddr; + bndPort = bPort; + return rep; } } // namespace aria2 diff --git a/src/SocksProxySocket.h b/src/SocksProxySocket.h index 93854c9d..f0979b45 100644 --- a/src/SocksProxySocket.h +++ b/src/SocksProxySocket.h @@ -58,8 +58,20 @@ private: // The socket is not shared because as it is used as some kind of controller, // when it is closed, all related proxy connections are closed. std::unique_ptr socket_; - std::vector bndAddrs_; - std::vector bndPorts_; + + enum SocksProxyCmd { + SOCKS_CMD_CONNECT = 1, + SOCKS_CMD_BIND = 2, + SOCKS_CMD_UDP_ASSOCIATE = 3, + }; + + // When allowEmtpy is true, allow dst* to be empty. + void sendCmd(SocksProxyCmd cmd, const std::string& dstAddr, uint16_t dstPort, + bool allowEmpty); + // Returns status replied from proxy server. 0 is OK. > 0 with error messages. + // < 0 for invalid replies from server. + // bndFamily may be AF_INET + AF_INET6 to indicate domain format. + int receiveReply(int& bndFamily, std::string& bndAddr, uint16_t& bndPort); public: SocksProxySocket(int family); @@ -75,10 +87,8 @@ public: // Negotiate authentication that returns selected auth method in 0-255. // When no auth method is selected, return -1. - // Auth methods in expected should be from enum SocksProxyAuthMethod, - // which has its auth handlers. // Returned auth method SHOULD be in expected but it is not checked. - int negotiateAuth(std::vector expected); + int negotiateAuth(std::vector expected); // Username/Password authentication. // user and pass should not be empty. @@ -87,17 +97,9 @@ public: // Create an UDP association to start UDP proxy. // Leave listen host and port empty / 0 to indicate no receiving from proxy. - // Returns -1 when error, otherwise the index to get the bnd addr and port. - // Set bndAddrPtr and bndPortPtr to directly get the result bnd addr and port. - ssize_t startUdpAssociate(const std::string& listenAddr, uint16_t listenPort, - std::pair bnd = {}); - - // Get bnd addr and port via index i. - // i is not checked and should be got from start*Proxy methods. - std::pair getBnd(size_t i) - { - return std::make_pair(bndAddrs_[i], bndPorts_[i]); - } + // Returns status replied from proxy server. 0 is OK. + int startUdpAssociate(const std::string& listenAddr, uint16_t listenPort, + std::string& bndAddr, uint16_t& bndPort); }; } // namespace aria2