Reduce the parsing to a single state machine, automatic removal of objects and destruction of sockets/streams

SOCKS.cpp

@@ -23,28 +23,33 @@ namespace proxy
 		}
 	}
 
-	void SOCKSHandler::Terminate() {
+	void SOCKSHandler::Done() {
-		CloseStream();
+		if (m_parent) m_parent->RemoveHandler (shared_from_this ());
-		CloseSock();
-		delete this; // HACK: ew
 	}
 
-	boost::asio::const_buffers_1 SOCKSHandler::GenerateSOCKS5SelectAuth(SOCKSHandler::authMethods method)
+	void SOCKSHandler::Terminate() {
-	{
+		if (dead.exchange(true)) return;
-		response[0] = '\x05'; //Version
+		if (m_sock) {
-		response[1] = method; //Response code
+			LogPrint(eLogDebug,"--- SOCKS close sock");
-		return boost::asio::const_buffers_1(response,2);
+			m_sock->close();
+			delete m_sock;
+			m_sock = nullptr;
+		}
+		if (m_stream) {
+			LogPrint(eLogDebug,"--- SOCKS close stream");
+			m_stream.reset ();
+		}
+		Done();
 	}
 
 	boost::asio::const_buffers_1 SOCKSHandler::GenerateSOCKS4Response(SOCKSHandler::errTypes error, uint32_t ip, uint16_t port)
 	{
 		assert(error >= SOCKS4_OK);
-		assert(type == ADDR_IPV4);
+		m_response[0] = '\x00'; //Version
-		response[0] = '\x00'; //Version
+		m_response[1] = error; //Response code
-		response[1] = error; //Response code
+		htobe16buf(m_response+2,port); //Port
-		htobe16buf(response+2,port); //Port
+		htobe32buf(m_response+4,ip); //IP
-		htobe32buf(response+4,ip); //IP
+		return boost::asio::const_buffers_1(m_response,8);
-		return boost::asio::const_buffers_1(response,8);
 	}
 
 	boost::asio::const_buffers_1 SOCKSHandler::GenerateSOCKS5Response(SOCKSHandler::errTypes error, SOCKSHandler::addrTypes type,
@@ -52,42 +57,43 @@ namespace proxy
 	{
 		size_t size;
 		assert(error <= SOCKS5_ADDR_UNSUP);
-		response[0] = '\x05'; //Version
+		m_response[0] = '\x05'; //Version
-		response[1] = error; //Response code
+		m_response[1] = error; //Response code
-		response[2] = '\x00'; //RSV
+		m_response[2] = '\x00'; //RSV
-		response[3] = type; //Address type
+		m_response[3] = type; //Address type
 		switch (type) {
 			case ADDR_IPV4:
 				size = 10;
-				htobe32buf(response+4,addr.ip);
+				htobe32buf(m_response+4,addr.ip);
 				break;
 			case ADDR_IPV6:
 				size = 22;
-				memcpy(response+4,addr.ipv6, 16);
+				memcpy(m_response+4,addr.ipv6, 16);
 				break;
 			case ADDR_DNS:
 				size = 7+addr.dns.size;
-				response[4] = addr.dns.size;
+				m_response[4] = addr.dns.size;
-				memcpy(response+5,addr.dns.value, addr.dns.size);
+				memcpy(m_response+5,addr.dns.value, addr.dns.size);
 				break;
 		}
-		htobe16buf(response+size-2,port); //Port
+		htobe16buf(m_response+size-2,port); //Port
-		return boost::asio::const_buffers_1(response,size);
+		return boost::asio::const_buffers_1(m_response,size);
 	}
 
-	/* Ah ah ah, you didn't say the magic word */
+	bool SOCKSHandler::Socks5ChooseAuth()
-	void SOCKSHandler::Socks5AuthNegoFailed()
 	{
+		m_response[0] = '\x05'; //Version
+		m_response[1] = m_authchosen; //Response code
+		boost::asio::const_buffers_1 response(m_response,2);
+		if (m_authchosen == AUTH_UNACCEPTABLE) {
 			LogPrint(eLogWarning,"--- SOCKS5 authentication negotiation failed");
-		boost::asio::async_write(*m_sock, GenerateSOCKS5SelectAuth(AUTH_UNACCEPTABLE),
+			boost::asio::async_write(*m_sock, response, std::bind(&SOCKSHandler::SentSocksFailed, this, std::placeholders::_1));
-					 std::bind(&SOCKSHandler::SentSocksFailed, this, std::placeholders::_1));
+			return false;
-	}
+		} else {
-
-	void SOCKSHandler::Socks5ChooseAuth()
-	{
 			LogPrint(eLogDebug,"--- SOCKS5 choosing authentication method: ", m_authchosen);
-		boost::asio::async_write(*m_sock, GenerateSOCKS5SelectAuth(m_authchosen),
+			boost::asio::async_write(*m_sock, response, std::bind(&SOCKSHandler::SentSocksResponse, this, std::placeholders::_1));
-					 std::bind(&SOCKSHandler::SentSocksResponse, this, std::placeholders::_1));
+			return true;
+		}
 	}
 
 	/* All hope is lost beyond this point */
@@ -123,114 +129,32 @@ namespace proxy
 				auto s = i2p::client::context.GetAddressBook().ToAddress(m_parent->GetLocalDestination()->GetIdentHash());
 				address ad; ad.dns.FromString(s);
 				//HACK only 16 bits passed in port as SOCKS5 doesn't allow for more
-				response = GenerateSOCKS5Response(SOCKS5_OK, ADDR_DNS, ad, m_stream->GetSendStreamID());
+				response = GenerateSOCKS5Response(SOCKS5_OK, ADDR_DNS, ad, m_stream->GetRecvStreamID());
 				break;
 		}
-		boost::asio::async_write(*m_sock, response, std::bind(&SOCKSHandler::SentSocksResponse, this, std::placeholders::_1));
+		boost::asio::async_write(*m_sock, response, std::bind(&SOCKSHandler::SentSocksDone, this, std::placeholders::_1));
-	}
-
-
-	void SOCKSHandler::CloseSock()
-	{
-		if (m_sock) {
-			LogPrint(eLogDebug,"--- SOCKS close sock");
-			m_sock->close();
-			delete m_sock;
-			m_sock = nullptr;
-		}
 	}
 
-       void SOCKSHandler::CloseStream()
+	void SOCKSHandler::EnterState(SOCKSHandler::state nstate, uint8_t parseleft) {
-       {
+		switch (nstate) {
-               if (m_stream) {
+			case GET_PORT: parseleft = 2; break;
-                       LogPrint(eLogDebug,"--- SOCKS close stream");
+			case GET_IPV4: m_addrtype = ADDR_IPV4; m_address.ip = 0; parseleft = 4; break;
-                       m_stream.reset ();
+			case GET4_IDENT: m_4aip = m_address.ip; break;
-               }
+			case GET4A_HOST:
-       }
+			case GET5_HOST: m_addrtype = ADDR_DNS; m_address.dns.size = 0; break;
-
+			case GET5_IPV6: m_addrtype = ADDR_IPV6; parseleft = 16; break;
-	std::size_t SOCKSHandler::HandleData(uint8_t *sock_buff, std::size_t len)
+			default:;
-	{
-		assert(len); // This should always be called with a least a byte left to parse
-		switch (m_state) {
-			case GET_VERSION:
-				return HandleVersion(sock_buff);
-			case SOCKS5_AUTHNEGO:
-				return HandleSOCKS5AuthNego(sock_buff,len);
-			case SOCKS_REQUEST:
-				return HandleSOCKSRequest(sock_buff,len);
-			default:
-				LogPrint(eLogError,"--- SOCKS state?? ", m_state);
-				Terminate();
-				return 0;
-		}
-	}
-
-	std::size_t SOCKSHandler::HandleVersion(uint8_t *sock_buff)
-	{
-		switch (*sock_buff) {
-			case SOCKS4:
-				m_state = SOCKS_REQUEST; // Switch to the 4 handler
-				m_pstate = GET_COMMAND; //Initialize the parser at the right position
-				break;
-			case SOCKS5:
-				m_state = SOCKS5_AUTHNEGO; // Switch to the 5 handler
-				m_pstate = GET5_AUTHNUM; //Initialize the parser at the right position
-				break;
-			default:
-				LogPrint(eLogError,"--- SOCKS rejected invalid version: ", ((int)*sock_buff));
-				Terminate();
-				return 0;
-		}
-		m_socksv = (SOCKSHandler::socksVersions) *sock_buff;
-		return 1;
-	}
-
-	std::size_t SOCKSHandler::HandleSOCKS5AuthNego(uint8_t *sock_buff, std::size_t len)
-	{
-		std::size_t rv = 0;
-		while (len > 0) {
-			rv++;
-			switch (m_pstate)
-			{
-				case GET5_AUTHNUM:
-					m_parseleft = *sock_buff;
-					m_pstate = GET5_AUTH;
-					break;
-				case GET5_AUTH:
-					m_parseleft --;
-					if (*sock_buff == AUTH_NONE)
-						m_authchosen = AUTH_NONE;
-					if ( m_parseleft == 0 ) {
-						if (m_authchosen == AUTH_UNACCEPTABLE) {
-							//TODO: we maybe want support for other methods!
-							LogPrint(eLogError,"--- SOCKS5 couldn't negotiate authentication");
-							Socks5AuthNegoFailed();
-							return 0;
-						}
-						m_pstate = GET5_REQUESTV;
-						m_state = SOCKS_REQUEST;
-						m_need_more = false;
-						Socks5ChooseAuth();
-						return rv;
-					}
-					break;
-				default:
-					LogPrint(eLogError,"--- SOCKS5 parse state?? ", m_pstate);
-					Terminate();
-					return 0;
-			}
-			sock_buff++;
-			len--;
 		}
-		return rv;
+		m_parseleft = parseleft;
+		m_state = nstate;
 	}
 
-	bool SOCKSHandler::ValidateSOCKSRequest() {
+	void SOCKSHandler::ValidateSOCKSRequest() {
 		if ( m_cmd != CMD_CONNECT ) {
 			//TODO: we need to support binds and other shit!
 			LogPrint(eLogError,"--- SOCKS unsupported command: ", m_cmd);
 			SocksRequestFailed(SOCKS5_CMD_UNSUP);
-			return false;
+			return;
 		}
 		//TODO: we may want to support other address types!
 		if ( m_addrtype != ADDR_DNS ) {
@@ -243,24 +167,48 @@ namespace proxy
 					break;
 			}
 			SocksRequestFailed(SOCKS5_ADDR_UNSUP);
-			return false;
+			return;
 		}
 		//TODO: we may want to support other domains
 		if(m_addrtype == ADDR_DNS && m_address.dns.ToString().find(".i2p") == std::string::npos) {
 			LogPrint(eLogError,"--- SOCKS invalid hostname: ", m_address.dns.ToString());
 			SocksRequestFailed(SOCKS5_ADDR_UNSUP);
-			return false;
+			return;
 		}
-		return true;
 	}
 
-	std::size_t SOCKSHandler::HandleSOCKSRequest(uint8_t *sock_buff, std::size_t len)
+	bool SOCKSHandler::HandleData(uint8_t *sock_buff, std::size_t len)
 	{
-		std::size_t rv = 0;
+		assert(len); // This should always be called with a least a byte left to parse
 		while (len > 0) {
-			rv++;
+			switch (m_state) {
-			switch (m_pstate)
+				case GET_SOCKSV:
-			{
+					m_socksv = (SOCKSHandler::socksVersions) *sock_buff;
+					switch (*sock_buff) {
+						case SOCKS4:
+							EnterState(GET_COMMAND); //Initialize the parser at the right position
+							break;
+						case SOCKS5:
+							EnterState(GET5_AUTHNUM); //Initialize the parser at the right position
+							break;
+						default:
+							LogPrint(eLogError,"--- SOCKS rejected invalid version: ", ((int)*sock_buff));
+							Terminate();
+							return false;
+					}
+					break;
+				case GET5_AUTHNUM:
+					EnterState(GET5_AUTH, *sock_buff);
+					break;
+				case GET5_AUTH:
+					m_parseleft --;
+					if (*sock_buff == AUTH_NONE)
+						m_authchosen = AUTH_NONE;
+					if ( m_parseleft == 0 ) {
+						if (!Socks5ChooseAuth()) return false;
+						EnterState(GET5_REQUESTV);
+					}
+					break;
 				case GET_COMMAND:
 					switch (*sock_buff) {
 						case CMD_CONNECT:
@@ -271,12 +219,12 @@ namespace proxy
 						default:
 							LogPrint(eLogError,"--- SOCKS invalid command: ", ((int)*sock_buff));
 							SocksRequestFailed(SOCKS5_GEN_FAIL);
-							return 0;
+							return false;
 					}
 					m_cmd = (SOCKSHandler::cmdTypes)*sock_buff;
 					switch (m_socksv) {
-						case SOCKS5: m_pstate = GET5_GETRSV; break;
+						case SOCKS5: EnterState(GET5_GETRSV); break;
-						case SOCKS4: m_pstate = GET_PORT;  m_parseleft = 2; break;
+						case SOCKS4: EnterState(GET_PORT); break;
 					}
 					break;
 				case GET_PORT:
@@ -284,8 +232,8 @@ namespace proxy
 					m_parseleft--;
 					if (m_parseleft == 0) {
 						switch (m_socksv) {
-							case SOCKS5: m_pstate = DONE; break;
+							case SOCKS5: EnterState(DONE); break;
-							case SOCKS4: m_pstate = GET_IPV4; m_address.ip = 0; m_parseleft = 4; break;
+							case SOCKS4: EnterState(GET_IPV4); break;
 						}
 					}
 					break;
@@ -294,32 +242,29 @@ namespace proxy
 					m_parseleft--;
 					if (m_parseleft == 0) {
 						switch (m_socksv) {
-							case SOCKS5: m_pstate = GET_PORT;  m_parseleft = 2; break;
+							case SOCKS5: EnterState(GET_PORT); break;
-							case SOCKS4: m_pstate = GET4_IDENT; m_4aip = m_address.ip; break;
+							case SOCKS4: EnterState(GET4_IDENT); m_4aip = m_address.ip; break;
 						}
 					}
 					break;
 				case GET4_IDENT:
 					if (!*sock_buff) {
 						if( m_4aip == 0 || m_4aip > 255 ) {
-							m_addrtype = ADDR_IPV4;
+							EnterState(DONE);
-							m_pstate = DONE;
 						} else {
-							m_addrtype = ADDR_DNS;
+							EnterState(GET4A_HOST);
-							m_pstate = GET4A_HOST;
-							m_address.dns.size = 0;
 						}
 					}
 					break;
 				case GET4A_HOST:
 					if (!*sock_buff) {
-						m_pstate = DONE;
+						EnterState(DONE);
 						break;
 					}
 					if (m_address.dns.size >= max_socks_hostname_size) {
 						LogPrint(eLogError,"--- SOCKS4a destination is too large");
 						SocksRequestFailed(SOCKS4_FAIL);
-						return 0;
+						return false;
 					}
 					m_address.dns.push_back(*sock_buff);
 					break;
@@ -327,65 +272,56 @@ namespace proxy
 					if (*sock_buff != SOCKS5) {
 						LogPrint(eLogError,"--- SOCKS5 rejected unknown request version: ", ((int)*sock_buff));
 						SocksRequestFailed(SOCKS5_GEN_FAIL);
-						return 0;
+						return false;
 					}
-					m_pstate = GET_COMMAND;
+					EnterState(GET_COMMAND);
 					break;
 				case GET5_GETRSV:
 					if ( *sock_buff != 0 ) {
 						LogPrint(eLogError,"--- SOCKS5 unknown reserved field: ", ((int)*sock_buff));
 						SocksRequestFailed(SOCKS5_GEN_FAIL);
-						return 0;
+						return false;
 					}
-					m_pstate = GET5_GETADDRTYPE;
+					EnterState(GET5_GETADDRTYPE);
 					break;
 				case GET5_GETADDRTYPE:
 					switch (*sock_buff) {
-						case ADDR_IPV4: m_pstate = GET_IPV4; m_address.ip = 0; m_parseleft = 4; break;
+						case ADDR_IPV4: EnterState(GET_IPV4); break;
-						case ADDR_IPV6: m_pstate = GET5_IPV6; m_parseleft = 16; break;
+						case ADDR_IPV6: EnterState(GET5_IPV6); break;
-						case ADDR_DNS : m_pstate = GET5_HOST_SIZE; break;
+						case ADDR_DNS : EnterState(GET5_HOST_SIZE); break;
 						default:
 							LogPrint(eLogError,"--- SOCKS5 unknown address type: ", ((int)*sock_buff));
 							SocksRequestFailed(SOCKS5_GEN_FAIL);
-							return 0;
+							return false;
 					}
-					m_addrtype = (SOCKSHandler::addrTypes)*sock_buff;
 					break;
 				case GET5_IPV6:
 					m_address.ipv6[16-m_parseleft] = *sock_buff;
 					m_parseleft--;
-					if (m_parseleft == 0) {
+					if (m_parseleft == 0) EnterState(GET_PORT);
-						m_pstate = GET_PORT;
-						m_parseleft = 2;
-					}
 					break;
 				case GET5_HOST_SIZE:
-					m_parseleft = *sock_buff;
+					EnterState(GET5_HOST, *sock_buff);
-					m_address.dns.size = 0;
-					m_pstate = GET5_HOST;
 					break;
 				case GET5_HOST:
 					m_address.dns.push_back(*sock_buff);
 					m_parseleft--;
-					if (m_parseleft == 0) {
+					if (m_parseleft == 0) EnterState(GET_PORT);
-						m_pstate = GET_PORT;
-						m_parseleft = 2;
-					}
 					break;
 				default:
-					LogPrint(eLogError,"--- SOCKS parse state?? ", m_pstate);
+					LogPrint(eLogError,"--- SOCKS parse state?? ", m_state);
 					Terminate();
-					return 0;
+					return false;
-			}
-			if (m_pstate == DONE) {
-				m_state = READY;
-				m_need_more = false;
-				return (ValidateSOCKSRequest() ? rv : 0);
 			}
 			sock_buff++;
 			len--;
+			if (len && m_state == DONE) {
+				LogPrint(eLogError,"--- SOCKS rejected because we can't handle extra data");
+				SocksRequestFailed(SOCKS5_GEN_FAIL);
+				return false;
 			}
-		return rv;
+		}
+		return true;
 	}
 
 	void SOCKSHandler::HandleSockRecv(const boost::system::error_code & ecode, std::size_t len)
@@ -397,34 +333,19 @@ namespace proxy
 			return;
 		}
 
-		std::size_t pos = 0;
+		if (HandleData(m_sock_buff, len)) {
-		m_need_more = true;
+			if (m_state == DONE) {
-		while (pos != len && m_state != READY && m_need_more) {
-			assert(pos < len); //We are overflowing the buffer otherwise
-			std::size_t rv = HandleData(m_sock_buff + pos, len - pos);
-			if (!rv) return; //Something went wrong die misserably
-			pos += rv;
-		}
-
-		assert(!(m_state == READY && m_need_more));
-
-		if (m_state == READY) {
 				LogPrint(eLogInfo,"--- SOCKS requested ", m_address.dns.ToString(), ":" , m_port);
-			if (pos != len) {
-				LogPrint(eLogError,"--- SOCKS rejected because we can't handle extra data");
-				SocksRequestFailed(SOCKS5_GEN_FAIL);
-				return ;
-			}
-
 				m_parent->GetLocalDestination ()->CreateStream (
 						std::bind (&SOCKSHandler::HandleStreamRequestComplete,
 						this, std::placeholders::_1), m_address.dns.ToString(), m_port);
-		} else if (m_need_more) {
+			} else {
-			LogPrint (eLogDebug,"--- SOCKS Need more data");
 				AsyncSockRead();
 			}
 		}
 
+	}
+
 	void SOCKSHandler::SentSocksFailed(const boost::system::error_code & ecode)
 	{
 		if (!ecode) {
@@ -435,21 +356,26 @@ namespace proxy
 		}
 	}
 
-	void SOCKSHandler::SentSocksResponse(const boost::system::error_code & ecode)
+	void SOCKSHandler::SentSocksDone(const boost::system::error_code & ecode)
 	{
 		if (!ecode) {
-			if(m_state == READY) {
+			if (dead.exchange(true)) return;
 			LogPrint (eLogInfo,"--- SOCKS New I2PTunnel connection");
 			auto connection = std::make_shared<i2p::client::I2PTunnelConnection>((i2p::client::I2PTunnel *)m_parent, m_sock, m_stream);
 			m_parent->AddConnection (connection);
 			connection->I2PConnect ();
-			} else {
+			Done();
-				LogPrint (eLogDebug,"--- SOCKS Go to next state");
-				AsyncSockRead();
-			}
 		}
 		else
 		{
+			LogPrint (eLogError,"--- SOCKS Closing socket after completion reply because: ", ecode.message ());
+			Terminate();
+		}
+	}
+
+	void SOCKSHandler::SentSocksResponse(const boost::system::error_code & ecode)
+	{
+		if (ecode) {
 			LogPrint (eLogError,"--- SOCKS Closing socket after sending reply because: ", ecode.message ());
 			Terminate();
 		}
@@ -458,7 +384,6 @@ namespace proxy
 	void SOCKSHandler::HandleStreamRequestComplete (std::shared_ptr<i2p::stream::Stream> stream)
 	{
 		if (stream) {
-			//TODO: test the stream is open if it isn't fail with connrefused
 			m_stream = stream;
 			SocksRequestSuccess();
 		} else {
@@ -467,7 +392,6 @@ namespace proxy
 		}
 	}
 
-
 	void SOCKSServer::Start ()
 	{
 		m_Acceptor.listen ();
@@ -479,6 +403,7 @@ namespace proxy
 		m_Acceptor.close();
 		m_Timer.cancel ();
 		ClearConnections ();
+		ClearHandlers();
 	}
 
 	void SOCKSServer::Accept ()
@@ -488,12 +413,29 @@ namespace proxy
 			std::placeholders::_1, newSocket));
 	}
 
+	void  SOCKSServer::AddHandler (std::shared_ptr<SOCKSHandler> handler) {
+		std::unique_lock<std::mutex> l(m_HandlersMutex);
+		m_Handlers.insert (handler);
+	}
+
+	void  SOCKSServer::RemoveHandler (std::shared_ptr<SOCKSHandler> handler)
+	{
+		std::unique_lock<std::mutex> l(m_HandlersMutex);
+		m_Handlers.erase (handler);
+	}
+
+	void  SOCKSServer::ClearHandlers ()
+	{
+		std::unique_lock<std::mutex> l(m_HandlersMutex);
+		m_Handlers.clear ();
+	}
+
 	void SOCKSServer::HandleAccept (const boost::system::error_code& ecode, boost::asio::ip::tcp::socket * socket)
 	{
 		if (!ecode)
 		{
 			LogPrint(eLogDebug,"--- SOCKS accepted");
-			new SOCKSHandler(this, socket);
+			AddHandler(std::make_shared<SOCKSHandler> (this, socket));
 			Accept();
 		}
 		else

SOCKS.h

@@ -3,7 +3,10 @@
 
 #include <memory>
 #include <string>
+#include <set>
 #include <boost/asio.hpp>
+#include <mutex>
+#include <atomic>
 #include "Identity.h"
 #include "Streaming.h"
 #include "I2PTunnel.h"
@@ -29,16 +32,10 @@ namespace proxy
 	};
 
 	class SOCKSServer;
-	class SOCKSHandler {
+	class SOCKSHandler: public std::enable_shared_from_this<SOCKSHandler> {
 		private:
 			enum state {
-				GET_VERSION, // Get SOCKS version
+				GET_SOCKSV,
-				SOCKS5_AUTHNEGO, //Authentication negotiation
-				SOCKS5_AUTH, //Authentication
-				SOCKS_REQUEST, //Request
-				READY // Ready to connect
-			};
-			enum parseState {
 				GET_COMMAND,
 				GET_PORT,
 				GET_IPV4,
@@ -95,25 +92,21 @@ namespace proxy
 				uint8_t ipv6[16];
 			};
 
-			std::size_t HandleData(uint8_t *sock_buff, std::size_t len);
+			void EnterState(state nstate, uint8_t parseleft = 1);
-			std::size_t HandleVersion(uint8_t *sock_buff);
+			bool HandleData(uint8_t *sock_buff, std::size_t len);
-			std::size_t HandleSOCKS5AuthNego(uint8_t *sock_buff, std::size_t len);
+			void ValidateSOCKSRequest();
-			std::size_t HandleSOCKSRequest(uint8_t *sock_buff, std::size_t len);
-			bool ValidateSOCKSRequest();
 			void HandleSockRecv(const boost::system::error_code & ecode, std::size_t bytes_transfered);
+			void Done();
 			void Terminate();
-			void CloseSock();
-			void CloseStream();
 			void AsyncSockRead();
 			boost::asio::const_buffers_1 GenerateSOCKS5SelectAuth(authMethods method);
 			boost::asio::const_buffers_1 GenerateSOCKS4Response(errTypes error, uint32_t ip, uint16_t port);
 			boost::asio::const_buffers_1 GenerateSOCKS5Response(errTypes error, addrTypes type, const address &addr, uint16_t port);
-			void Socks5AuthNegoFailed();
+			bool Socks5ChooseAuth();
-			void Socks5ChooseAuth();
 			void SocksRequestFailed(errTypes error);
 			void SocksRequestSuccess();
-			//HACK: we need to pass the shared_ptr to ensure the buffer will live enough
 			void SentSocksFailed(const boost::system::error_code & ecode);
+			void SentSocksDone(const boost::system::error_code & ecode);
 			void SentSocksResponse(const boost::system::error_code & ecode);
 			void HandleStreamRequestComplete (std::shared_ptr<i2p::stream::Stream> stream);
 
@@ -121,9 +114,7 @@ namespace proxy
 			SOCKSServer * m_parent;
 			boost::asio::ip::tcp::socket * m_sock;
 			std::shared_ptr<i2p::stream::Stream> m_stream;
-			state m_state;
+			uint8_t m_response[7+max_socks_hostname_size];
-			parseState m_pstate;
-			uint8_t response[7+max_socks_hostname_size];
 			address m_address; //Address
 			uint32_t m_4aip; //Used in 4a requests
 			uint16_t m_port;
@@ -133,18 +124,30 @@ namespace proxy
 			addrTypes m_addrtype; //Address type chosen
 			socksVersions m_socksv; //Socks version
 			cmdTypes m_cmd; // Command requested
-			bool m_need_more; //The parser still needs to receive more data
+			state m_state;
+			std::atomic<bool> dead; //To avoid cleaning up multiple times
 
 		public:
 			SOCKSHandler(SOCKSServer * parent, boost::asio::ip::tcp::socket * sock) : 
-				m_parent(parent), m_sock(sock), m_stream(nullptr), m_state(GET_VERSION),
+				m_parent(parent), m_sock(sock), m_stream(nullptr),
-				m_authchosen(AUTH_UNACCEPTABLE), m_addrtype(ADDR_IPV4)
+				m_authchosen(AUTH_UNACCEPTABLE), m_addrtype(ADDR_IPV4), dead(false)
-				{ m_address.ip = 0; AsyncSockRead(); }
+				{ m_address.ip = 0; AsyncSockRead(); EnterState(GET_SOCKSV); }
-			~SOCKSHandler() { CloseSock(); CloseStream(); }
+			~SOCKSHandler() { Terminate(); }
 	};
 
 	class SOCKSServer: public i2p::client::I2PTunnel
 	{
+		private:
+			std::set<std::shared_ptr<SOCKSHandler> > m_Handlers;
+			boost::asio::ip::tcp::acceptor m_Acceptor;
+			boost::asio::deadline_timer m_Timer;
+			std::mutex m_HandlersMutex;
+
+		private:
+
+			void Accept();
+			void HandleAccept(const boost::system::error_code& ecode, boost::asio::ip::tcp::socket * socket);
+
 		public:
 			SOCKSServer(int port) : I2PTunnel(nullptr),
 				m_Acceptor (GetService (), boost::asio::ip::tcp::endpoint (boost::asio::ip::tcp::v4(), port)),
@@ -153,16 +156,9 @@ namespace proxy
 
 			void Start ();
 			void Stop ();
-
+			void AddHandler (std::shared_ptr<SOCKSHandler> handler);
-		private:
+			void RemoveHandler (std::shared_ptr<SOCKSHandler> handler);
-
+			void ClearHandlers ();
-			void Accept();
-			void HandleAccept(const boost::system::error_code& ecode, boost::asio::ip::tcp::socket * socket);
-
-		private:
-
-			boost::asio::ip::tcp::acceptor m_Acceptor;
-			boost::asio::deadline_timer m_Timer;
 	};
 
 	typedef SOCKSServer SOCKSProxy;